Line data    Source code

       1             : /*
       2             :  *
       3             :  *    Copyright (c) 2020 Project CHIP Authors
       4             :  *
       5             :  *    Licensed under the Apache License, Version 2.0 (the "License");
       6             :  *    you may not use this file except in compliance with the License.
       7             :  *    You may obtain a copy of the License at
       8             :  *
       9             :  *        http://www.apache.org/licenses/LICENSE-2.0
      10             :  *
      11             :  *    Unless required by applicable law or agreed to in writing, software
      12             :  *    distributed under the License is distributed on an "AS IS" BASIS,
      13             :  *    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
      14             :  *    See the License for the specific language governing permissions and
      15             :  *    limitations under the License.
      16             :  */
      17             : 
      18             : #pragma once
      19             : 
      20             : #include <system/SystemPacketBuffer.h>
      21             : 
      22             : #include <lib/dnssd/minimal_mdns/Parser.h>
      23             : #include <lib/dnssd/minimal_mdns/core/DnsHeader.h>
      24             : #include <lib/dnssd/minimal_mdns/records/ResourceRecord.h>
      25             : 
      26             : namespace mdns {
      27             : namespace Minimal {
      28             : 
      29             : /// Writes a MDNS reply into a given packet buffer.
      30             : class ResponseBuilder
      31             : {
      32             : public:
      33           2 :     ResponseBuilder() : mHeader(nullptr), mEndianOutput(nullptr, 0), mWriter(&mEndianOutput) {}
      34             :     ResponseBuilder(chip::System::PacketBufferHandle && packet) :
      35             :         mHeader(nullptr), mEndianOutput(nullptr, 0), mWriter(&mEndianOutput)
      36             :     {
      37             :         Reset(std::move(packet));
      38             :     }
      39             : 
      40         110 :     ResponseBuilder & Reset(chip::System::PacketBufferHandle && packet)
      41             :     {
      42         110 :         mPacket = std::move(packet);
      43         110 :         mHeader = HeaderRef(mPacket->Start());
      44             : 
      45         110 :         if (mPacket->AvailableDataLength() >= HeaderRef::kSizeBytes)
      46             :         {
      47         110 :             mPacket->SetDataLength(HeaderRef::kSizeBytes);
      48         110 :             mHeader.Clear();
      49         110 :             mBuildOk = true;
      50             :         }
      51             :         else
      52             :         {
      53           0 :             mBuildOk = false;
      54             :         }
      55             : 
      56         110 :         mHeader.SetFlags(mHeader.GetFlags().SetResponse().SetAuthoritative());
      57             : 
      58         110 :         mEndianOutput =
      59         110 :             chip::Encoding::BigEndian::BufferWriter(mPacket->Start(), mPacket->DataLength() + mPacket->AvailableDataLength());
      60         110 :         mEndianOutput.Skip(mPacket->DataLength());
      61             : 
      62         110 :         mWriter.Reset();
      63             : 
      64         110 :         return *this;
      65             :     }
      66             : 
      67             :     CHECK_RETURN_VALUE
      68         110 :     chip::System::PacketBufferHandle ReleasePacket()
      69             :     {
      70         110 :         mHeader  = HeaderRef(nullptr);
      71         110 :         mBuildOk = false;
      72         110 :         return std::move(mPacket);
      73             :     }
      74             : 
      75         110 :     bool HasResponseRecords() const
      76             :     {
      77         110 :         return (mHeader.GetAnswerCount() != 0) || (mHeader.GetAuthorityCount() != 0) || (mHeader.GetAdditionalCount() != 0);
      78             :     }
      79             : 
      80         130 :     HeaderRef & Header() { return mHeader; }
      81             : 
      82             :     /// Attempts to add a record to the currentsystem packet buffer.
      83             :     /// On success, the packet buffer data length is updated.
      84             :     /// On failure, the packet buffer data length is NOT updated and header is unchanged.
      85         831 :     ResponseBuilder & AddRecord(ResourceType type, const ResourceRecord & record)
      86             :     {
      87         831 :         if (!mBuildOk)
      88             :         {
      89           0 :             return *this;
      90             :         }
      91             : 
      92         831 :         if (!record.Append(mHeader, type, mWriter))
      93             :         {
      94          10 :             mBuildOk = false;
      95             :         }
      96             :         else
      97             :         {
      98         821 :             VerifyOrDie(mEndianOutput.Fit()); // should be guaranteed because record Append succeeded
      99         821 :             mPacket->SetDataLength(static_cast<uint16_t>(mEndianOutput.Needed()));
     100             :         }
     101             : 
     102         831 :         return *this;
     103             :     }
     104             : 
     105          32 :     ResponseBuilder & AddQuery(const QueryData & query)
     106             :     {
     107          32 :         if (!mBuildOk)
     108             :         {
     109           0 :             return *this;
     110             :         }
     111             : 
     112          32 :         if (!query.Append(mHeader, mWriter))
     113             :         {
     114           0 :             mBuildOk = false;
     115             :         }
     116             :         else
     117             :         {
     118          32 :             mPacket->SetDataLength(static_cast<uint16_t>(mEndianOutput.Needed()));
     119             :         }
     120          32 :         return *this;
     121             :     }
     122             : 
     123        1662 :     bool Ok() const { return mBuildOk; }
     124         980 :     bool HasPacketBuffer() const { return !mPacket.IsNull(); }
     125             : 
     126             : private:
     127             :     chip::System::PacketBufferHandle mPacket;
     128             :     HeaderRef mHeader;
     129             :     chip::Encoding::BigEndian::BufferWriter mEndianOutput;
     130             :     RecordWriter mWriter;
     131             :     bool mBuildOk = false;
     132             : };
     133             : 
     134             : } // namespace Minimal
     135             : } // namespace mdns