github.com/slackhq/nebula@v1.9.0/udp/udp_tester.go (about)

     1  //go:build e2e_testing
     2  // +build e2e_testing
     3  
     4  package udp
     5  
     6  import (
     7  	"fmt"
     8  	"io"
     9  	"net"
    10  	"sync/atomic"
    11  
    12  	"github.com/sirupsen/logrus"
    13  	"github.com/slackhq/nebula/config"
    14  	"github.com/slackhq/nebula/firewall"
    15  	"github.com/slackhq/nebula/header"
    16  )
    17  
    18  type Packet struct {
    19  	ToIp     net.IP
    20  	ToPort   uint16
    21  	FromIp   net.IP
    22  	FromPort uint16
    23  	Data     []byte
    24  }
    25  
    26  func (u *Packet) Copy() *Packet {
    27  	n := &Packet{
    28  		ToIp:     make(net.IP, len(u.ToIp)),
    29  		ToPort:   u.ToPort,
    30  		FromIp:   make(net.IP, len(u.FromIp)),
    31  		FromPort: u.FromPort,
    32  		Data:     make([]byte, len(u.Data)),
    33  	}
    34  
    35  	copy(n.ToIp, u.ToIp)
    36  	copy(n.FromIp, u.FromIp)
    37  	copy(n.Data, u.Data)
    38  	return n
    39  }
    40  
    41  type TesterConn struct {
    42  	Addr *Addr
    43  
    44  	RxPackets chan *Packet // Packets to receive into nebula
    45  	TxPackets chan *Packet // Packets transmitted outside by nebula
    46  
    47  	closed atomic.Bool
    48  	l      *logrus.Logger
    49  }
    50  
    51  func NewListener(l *logrus.Logger, ip net.IP, port int, _ bool, _ int) (Conn, error) {
    52  	return &TesterConn{
    53  		Addr:      &Addr{ip, uint16(port)},
    54  		RxPackets: make(chan *Packet, 10),
    55  		TxPackets: make(chan *Packet, 10),
    56  		l:         l,
    57  	}, nil
    58  }
    59  
    60  // Send will place a UdpPacket onto the receive queue for nebula to consume
    61  // this is an encrypted packet or a handshake message in most cases
    62  // packets were transmitted from another nebula node, you can send them with Tun.Send
    63  func (u *TesterConn) Send(packet *Packet) {
    64  	if u.closed.Load() {
    65  		return
    66  	}
    67  
    68  	h := &header.H{}
    69  	if err := h.Parse(packet.Data); err != nil {
    70  		panic(err)
    71  	}
    72  	if u.l.Level >= logrus.DebugLevel {
    73  		u.l.WithField("header", h).
    74  			WithField("udpAddr", fmt.Sprintf("%v:%v", packet.FromIp, packet.FromPort)).
    75  			WithField("dataLen", len(packet.Data)).
    76  			Debug("UDP receiving injected packet")
    77  	}
    78  	u.RxPackets <- packet
    79  }
    80  
    81  // Get will pull a UdpPacket from the transmit queue
    82  // nebula meant to send this message on the network, it will be encrypted
    83  // packets were ingested from the tun side (in most cases), you can send them with Tun.Send
    84  func (u *TesterConn) Get(block bool) *Packet {
    85  	if block {
    86  		return <-u.TxPackets
    87  	}
    88  
    89  	select {
    90  	case p := <-u.TxPackets:
    91  		return p
    92  	default:
    93  		return nil
    94  	}
    95  }
    96  
    97  //********************************************************************************************************************//
    98  // Below this is boilerplate implementation to make nebula actually work
    99  //********************************************************************************************************************//
   100  
   101  func (u *TesterConn) WriteTo(b []byte, addr *Addr) error {
   102  	if u.closed.Load() {
   103  		return io.ErrClosedPipe
   104  	}
   105  
   106  	p := &Packet{
   107  		Data:     make([]byte, len(b), len(b)),
   108  		FromIp:   make([]byte, 16),
   109  		FromPort: u.Addr.Port,
   110  		ToIp:     make([]byte, 16),
   111  		ToPort:   addr.Port,
   112  	}
   113  
   114  	copy(p.Data, b)
   115  	copy(p.ToIp, addr.IP.To16())
   116  	copy(p.FromIp, u.Addr.IP.To16())
   117  
   118  	u.TxPackets <- p
   119  	return nil
   120  }
   121  
   122  func (u *TesterConn) ListenOut(r EncReader, lhf LightHouseHandlerFunc, cache *firewall.ConntrackCacheTicker, q int) {
   123  	plaintext := make([]byte, MTU)
   124  	h := &header.H{}
   125  	fwPacket := &firewall.Packet{}
   126  	ua := &Addr{IP: make([]byte, 16)}
   127  	nb := make([]byte, 12, 12)
   128  
   129  	for {
   130  		p, ok := <-u.RxPackets
   131  		if !ok {
   132  			return
   133  		}
   134  		ua.Port = p.FromPort
   135  		copy(ua.IP, p.FromIp.To16())
   136  		r(ua, plaintext[:0], p.Data, h, fwPacket, lhf, nb, q, cache.Get(u.l))
   137  	}
   138  }
   139  
   140  func (u *TesterConn) ReloadConfig(*config.C) {}
   141  
   142  func NewUDPStatsEmitter(_ []Conn) func() {
   143  	// No UDP stats for non-linux
   144  	return func() {}
   145  }
   146  
   147  func (u *TesterConn) LocalAddr() (*Addr, error) {
   148  	return u.Addr, nil
   149  }
   150  
   151  func (u *TesterConn) Rebind() error {
   152  	return nil
   153  }
   154  
   155  func (u *TesterConn) Close() error {
   156  	if u.closed.CompareAndSwap(false, true) {
   157  		close(u.RxPackets)
   158  		close(u.TxPackets)
   159  	}
   160  	return nil
   161  }