github.com/gopacket/gopacket@v1.1.0/defrag/lcmdefrag/lcmdefrag.go (about)

     1  // Copyright 2018 Google, Inc. All rights reserved.
     2  //
     3  // Use of this source code is governed by a BSD-style license
     4  // that can be found in the LICENSE file in the root of the source
     5  // tree.
     6  
     7  // Package lcmdefrag contains a defragmenter for LCM messages.
     8  package lcmdefrag
     9  
    10  import (
    11  	"fmt"
    12  	"time"
    13  
    14  	"github.com/gopacket/gopacket"
    15  	"github.com/gopacket/gopacket/layers"
    16  )
    17  
    18  const (
    19  	// Packages are cleaned up/removed after no input was received for this
    20  	// amount of seconds.
    21  	timeout time.Duration = 3 * time.Second
    22  )
    23  
    24  type lcmPacket struct {
    25  	lastPacket time.Time
    26  	done       bool
    27  	recFrags   uint16
    28  	totalFrags uint16
    29  	frags      map[uint16]*layers.LCM
    30  }
    31  
    32  // LCMDefragmenter supports defragmentation of LCM messages.
    33  //
    34  // References
    35  //
    36  //	https://lcm-proj.github.io/
    37  //	https://github.com/lcm-proj/lcm
    38  type LCMDefragmenter struct {
    39  	packets map[uint32]*lcmPacket
    40  }
    41  
    42  func newLCMPacket(totalFrags uint16) *lcmPacket {
    43  	return &lcmPacket{
    44  		done:       false,
    45  		recFrags:   0,
    46  		totalFrags: totalFrags,
    47  		frags:      make(map[uint16]*layers.LCM),
    48  	}
    49  }
    50  
    51  // NewLCMDefragmenter returns a new LCMDefragmenter.
    52  func NewLCMDefragmenter() *LCMDefragmenter {
    53  	return &LCMDefragmenter{
    54  		packets: make(map[uint32]*lcmPacket),
    55  	}
    56  }
    57  
    58  func (lp *lcmPacket) append(in *layers.LCM) {
    59  	lp.frags[in.FragmentNumber] = in
    60  	lp.recFrags++
    61  	lp.lastPacket = time.Now()
    62  }
    63  
    64  func (lp *lcmPacket) assemble() (out *layers.LCM, err error) {
    65  	var blob []byte
    66  
    67  	//Extract packets
    68  	for i := uint16(0); i < lp.totalFrags; i++ {
    69  		fragment, ok := lp.frags[i]
    70  		if !ok {
    71  			err = fmt.Errorf("Tried to defragment incomplete packet. Waiting "+
    72  				"for more potential (unordered) packets... %d", i)
    73  			return
    74  		}
    75  
    76  		// For the very first packet, we also want the header.
    77  		if i == 0 {
    78  			blob = append(blob, fragment.LayerContents()...)
    79  		}
    80  
    81  		// Append the data for each packet.
    82  		blob = append(blob, fragment.Payload()...)
    83  	}
    84  
    85  	packet := gopacket.NewPacket(blob, layers.LayerTypeLCM, gopacket.NoCopy)
    86  	lcmHdrLayer := packet.Layer(layers.LayerTypeLCM)
    87  	out, ok := lcmHdrLayer.(*layers.LCM)
    88  	if !ok {
    89  		err = fmt.Errorf("Error while decoding the defragmented packet. " +
    90  			"Erasing/dropping packet.")
    91  	}
    92  
    93  	lp.done = true
    94  
    95  	return
    96  }
    97  
    98  func (ld *LCMDefragmenter) cleanUp() {
    99  	for key, packet := range ld.packets {
   100  		if packet.done || time.Now().Sub(packet.lastPacket) > timeout {
   101  			delete(ld.packets, key)
   102  		}
   103  	}
   104  }
   105  
   106  // Defrag takes a reference to an LCM packet and processes it.
   107  // In case the packet does not need to be defragmented, it immediately returns
   108  // the as in passed reference. In case in was the last missing fragment, out
   109  // will be the defragmented packet. If in was a fragment, but we are awaiting
   110  // more, out will be set to nil.
   111  // In the case that in was nil, we will just run the internal cleanup of the
   112  // defragmenter that times out packages.
   113  // If an error was encountered during defragmentation, out will also be nil,
   114  // while err will contain further information on the failure.
   115  func (ld *LCMDefragmenter) Defrag(in *layers.LCM) (out *layers.LCM, err error) {
   116  	// Timeout old packages and erase error prone ones.
   117  	ld.cleanUp()
   118  
   119  	// For running cleanup only
   120  	if in == nil {
   121  		return
   122  	}
   123  
   124  	// Quick check if this is acutally a single packet. In that case, just
   125  	// return it quickly.
   126  	if !in.Fragmented {
   127  		out = in
   128  		return
   129  	}
   130  
   131  	// Do we need to start a new fragments obj?
   132  	if _, ok := ld.packets[in.SequenceNumber]; !ok {
   133  		ld.packets[in.SequenceNumber] = newLCMPacket(in.TotalFragments)
   134  	}
   135  
   136  	// Append the packet
   137  	ld.packets[in.SequenceNumber].append(in)
   138  
   139  	// Check if this is the last package of that series
   140  	if ld.packets[in.SequenceNumber].recFrags == in.TotalFragments {
   141  		out, err = ld.packets[in.SequenceNumber].assemble()
   142  	}
   143  
   144  	return
   145  }