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            3 :     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          104 :     ResponseBuilder & Reset(chip::System::PacketBufferHandle && packet)
      41              :     {
      42          104 :         mPacket = std::move(packet);
      43          104 :         mHeader = HeaderRef(mPacket->Start());
      44              : 
      45          104 :         if (mPacket->AvailableDataLength() >= HeaderRef::kSizeBytes)
      46              :         {
      47          104 :             mPacket->SetDataLength(HeaderRef::kSizeBytes);
      48          104 :             mHeader.Clear();
      49          104 :             mBuildOk = true;
      50              :         }
      51              :         else
      52              :         {
      53            0 :             mBuildOk = false;
      54              :         }
      55              : 
      56          104 :         mHeader.SetFlags(mHeader.GetFlags().SetResponse().SetAuthoritative());
      57              : 
      58          104 :         mEndianOutput =
      59          104 :             chip::Encoding::BigEndian::BufferWriter(mPacket->Start(), mPacket->DataLength() + mPacket->AvailableDataLength());
      60          104 :         mEndianOutput.Skip(mPacket->DataLength());
      61              : 
      62          104 :         mWriter.Reset();
      63              : 
      64          104 :         return *this;
      65              :     }
      66              : 
      67              :     CHECK_RETURN_VALUE
      68          104 :     chip::System::PacketBufferHandle ReleasePacket()
      69              :     {
      70          104 :         mHeader  = HeaderRef(nullptr);
      71          104 :         mBuildOk = false;
      72          104 :         return std::move(mPacket);
      73              :     }
      74              : 
      75          104 :     bool HasResponseRecords() const
      76              :     {
      77          104 :         return (mHeader.GetAnswerCount() != 0) || (mHeader.GetAuthorityCount() != 0) || (mHeader.GetAdditionalCount() != 0);
      78              :     }
      79              : 
      80          124 :     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          775 :     ResponseBuilder & AddRecord(ResourceType type, const ResourceRecord & record)
      86              :     {
      87          775 :         if (!mBuildOk)
      88              :         {
      89            0 :             return *this;
      90              :         }
      91              : 
      92          775 :         if (!record.Append(mHeader, type, mWriter))
      93              :         {
      94           10 :             mBuildOk = false;
      95              :         }
      96              :         else
      97              :         {
      98          765 :             VerifyOrDie(mEndianOutput.Fit()); // should be guaranteed because record Append succeeded
      99          765 :             mPacket->SetDataLength(static_cast<uint16_t>(mEndianOutput.Needed()));
     100              :         }
     101              : 
     102          775 :         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         1550 :     bool Ok() const { return mBuildOk; }
     124          910 :     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