github.com/aporeto-inc/trireme-lib@v10.358.0+incompatible/controller/pkg/flowtracking/flowtracking.go (about)

     1  // +build linux
     2  
     3  package flowtracking
     4  
     5  import (
     6  	"context"
     7  	"fmt"
     8  	"net"
     9  
    10  	"github.com/mdlayher/netlink"
    11  	"github.com/ti-mo/conntrack"
    12  )
    13  
    14  // Client is a flow update client
    15  type Client struct {
    16  	conn *conntrack.Conn
    17  }
    18  
    19  // NewClient creates a new flow tracking client. s
    20  func NewClient(ctx context.Context) (*Client, error) {
    21  	c, err := conntrack.Dial(&netlink.Config{
    22  		// Enable this when the netlink PR is merged.
    23  		DisableNSLockThread: true,
    24  	})
    25  	if err != nil {
    26  		return nil, fmt.Errorf("flow tracker is unable to dial netlink: %s", err)
    27  	}
    28  
    29  	client := &Client{conn: c}
    30  	go func() {
    31  		<-ctx.Done()
    32  		client.conn.Close() // nolint errcheck
    33  	}()
    34  
    35  	return client, nil
    36  }
    37  
    38  // Close will close the connection of the client.
    39  func (c *Client) Close() error {
    40  	return c.conn.Close()
    41  }
    42  
    43  // UpdateMark updates the mark of the flow. Caller must indicate if this is an application
    44  // flow or a network flow.
    45  func (c *Client) UpdateMark(ipSrc, ipDst net.IP, protonum uint8, srcport, dstport uint16, newmark uint32, network bool) error {
    46  
    47  	if network {
    48  		return c.UpdateNetworkFlowMark(ipSrc, ipDst, protonum, srcport, dstport, newmark)
    49  	}
    50  
    51  	return c.UpdateApplicationFlowMark(ipSrc, ipDst, protonum, srcport, dstport, newmark)
    52  }
    53  
    54  // GetOriginalDest gets the original destination ip, port and the mark on the packet
    55  func (c *Client) GetOriginalDest(ipSrc, ipDst net.IP, srcport, dstport uint16, protonum uint8) (net.IP, uint16, uint32, error) {
    56  
    57  	flow := conntrack.NewFlow(protonum, 0, ipSrc, ipDst, srcport, dstport, 0, 0)
    58  	origFlow, err := c.conn.Get(flow)
    59  	return origFlow.TupleOrig.IP.DestinationAddress, origFlow.TupleOrig.Proto.DestinationPort, origFlow.Mark, err
    60  }
    61  
    62  // UpdateNetworkFlowMark will update the mark for a flow based on packet information received
    63  // from the network. It will use the reverse tables in conntrack for that.
    64  func (c *Client) UpdateNetworkFlowMark(ipSrc, ipDst net.IP, protonum uint8, srcport, dstport uint16, newmark uint32) error {
    65  
    66  	f := newReplyFlow(protonum, 0, ipSrc, ipDst, srcport, dstport, 0, newmark)
    67  
    68  	return c.conn.Update(f)
    69  }
    70  
    71  // UpdateApplicationFlowMark will update the mark for a flow based on the packet information
    72  // received from an application. It will use the forward entries of conntrack for that.
    73  func (c *Client) UpdateApplicationFlowMark(ipSrc, ipDst net.IP, protonum uint8, srcport, dstport uint16, newmark uint32) error {
    74  
    75  	f := conntrack.NewFlow(protonum, 0, ipSrc, ipDst, srcport, dstport, 0, newmark)
    76  
    77  	return c.conn.Update(f)
    78  }
    79  
    80  // newReplyFlow will create a flow based on the reply tuple only. This will help us
    81  // update the mark without requiring knowledge of nats.
    82  func newReplyFlow(proto uint8, status conntrack.StatusFlag, srcAddr, destAddr net.IP, srcPort, destPort uint16, timeout, mark uint32) conntrack.Flow {
    83  
    84  	var f conntrack.Flow
    85  
    86  	f.Status.Value = status
    87  
    88  	f.Timeout = timeout
    89  	f.Mark = mark
    90  
    91  	// Set up TupleReply with source and destination inverted
    92  	f.TupleReply.IP.SourceAddress = srcAddr
    93  	f.TupleReply.IP.DestinationAddress = destAddr
    94  	f.TupleReply.Proto.SourcePort = srcPort
    95  	f.TupleReply.Proto.DestinationPort = destPort
    96  	f.TupleReply.Proto.Protocol = proto
    97  
    98  	return f
    99  }