device.go

// Copyright 2018 Dan Jacques. All rights reserved.
// Use of this source code is governed under the MIT License
// that can be found in the LICENSE file.

package proxy

import (
	"fmt"
	"net"
	"sync"
	"time"

	"github.com/danjacques/gopushpixels/device"
	"github.com/danjacques/gopushpixels/protocol"
	"github.com/danjacques/gopushpixels/support/bufferpool"
	"github.com/danjacques/gopushpixels/support/byteslicereader"
	"github.com/danjacques/gopushpixels/support/fmtutil"
	"github.com/danjacques/gopushpixels/support/logging"

	"github.com/pkg/errors"
	"github.com/prometheus/client_golang/prometheus"
)

// Device is a device.D that proxies for another device. Data written to this
// device can be sent to the proxied base device and/or captured, enabling
// man-in-the-middle operations on devices.
//
// Device should not be instantiated directly; instead, it should be obtained
// through a Manager.
type Device struct {
	// m is this device's proxy manager.
	m *Manager
	// logger is the resolved logger instance.
	logger logging.L

	// proxy is the local (proxy) device that the proxy endpoint is using.
	proxy device.Local
	// proxyHWAddr is the hardware address assigned to this proxy device. This is
	// a fake address, generated by AddressRegistry.
	proxyHWAddr net.HardwareAddr

	// monitoring is the device monitoring state.
	monitoring device.Monitoring

	// base is the base device that Device is proxying for.
	base device.D
	// baseID is base's cached ID.
	baseID string

	// createdTime is the time when this device was created.
	createdTime time.Time

	// doneC is this Device's done channel. It will close when either the device
	// has closed (explicit) or the device's base has closed (see
	// closeWhenBaseCloses()).
	doneC        chan struct{}
	shutdownOnce sync.Once

	basePacketsC     chan *packetData
	listenerPacketsC chan *packetData
	wg               sync.WaitGroup

	infoMu sync.Mutex
	info   device.Info

	counterLabels prometheus.Labels
}

// makeProxyDevice allocates a new proxy device with the assigned address,
// hwAddr.
//
// The device will proxy received data to d, and leverage d's discovery headers
// when constructing its own.
//
// Any packets received will be sent to l, the set of registered listeners for
// this device.
//
// The proxy device will close and terminate once the supplied Context is
// closed.
func makeProxyDevice(m *Manager, d device.D, hwAddr net.HardwareAddr) (*Device, error) {
	const chanSize = 1024

	// Create a PacketReader for our base device.
	dh := d.DiscoveryHeaders()
	basePR, err := dh.PacketReader()
	if err != nil {
		return nil, errors.Wrap(err, "could not create PacketReader")
	}

	// Open a listening socket for this device.
	conn, err := net.ListenUDP("udp4", &net.UDPAddr{IP: m.ProxyAddr})
	if err != nil {
		return nil, err
	}
	defer func() {
		if conn != nil {
			_ = conn.Close()
		}
	}()

	// Create a monitored Sender.
	s, err := d.Sender()
	if err != nil {
		return nil, err
	}

	pd := Device{
		m:      m,
		logger: logging.Must(m.Logger),

		proxy: device.Local{
			DeviceID:      hwAddr.String(),
			UDPPacketPool: m.udpPacketPool,
			Logger:        m.Logger,
		},
		proxyHWAddr: hwAddr,

		base:   d,
		baseID: d.ID(),

		createdTime: time.Now(),

		doneC:            make(chan struct{}),
		basePacketsC:     make(chan *packetData, chanSize),
		listenerPacketsC: make(chan *packetData, chanSize),
	}

	// Connect our callback and start our local listener.
	pd.proxy.OnPacketData = pd.onLocalPacketData
	pd.proxy.Start(conn)
	conn = nil // Owned by pd.proxy.

	// Generate/cache our metric labels for this device.
	pd.counterLabels = prometheus.Labels{
		"proxy_id": pd.proxy.DeviceID,
		"base_id":  pd.baseID,
	}

	// Watch base and close the proxy when it closes.
	go pd.closeWhenBaseCloses()

	// Process packets from the local Proxy connection and dispatch them to our
	// base device.
	//
	// This passes ownership of pw to the goroutine.
	pd.wg.Add(1)
	go func() {
		defer pd.wg.Done()
		pd.forwardPacketsToBase(s)
	}()

	// Process packets from the local Proxy connection and dispatch them to our
	// listeners.
	//
	// We do this independently from Router dispatching in order to buffer router
	// throughput from immediate Listener hiccups.
	pd.wg.Add(1)
	go func() {
		defer pd.wg.Done()
		pd.forwardPacketsToListeners(basePR)
	}()

	// Update our device monitoring.
	pd.monitoring.Update(&pd)

	return &pd, nil
}

func (pd *Device) String() string {
	return fmt.Sprintf("proxy.Device{%s proxying %s}", pd.proxy.DeviceID, pd.baseID)
}

// ID implements device.D.
func (pd *Device) ID() string { return pd.proxy.DeviceID }

// Ordinal implements device.D.
func (pd *Device) Ordinal() device.Ordinal {
	o := pd.base.Ordinal()
	o.Group += int(pd.m.GroupOffset)
	return o
}

// Sender implements device.D.
func (pd *Device) Sender() (device.Sender, error) { return pd.base.Sender() }

// DiscoveryHeaders implements device.D.
func (pd *Device) DiscoveryHeaders() *protocol.DiscoveryHeaders {
	proxyAddr := pd.proxy.Addr().(*net.UDPAddr)

	// Permute the current base proxy headers.
	dh := pd.base.DiscoveryHeaders().Clone()
	dh.SetIP4Address(proxyAddr.IP)
	dh.SetHardwareAddr(pd.proxyHWAddr)

	if dh.PixelPusher != nil {
		ordinal := pd.Ordinal()
		dh.PixelPusher.MyPort = uint16(proxyAddr.Port)
		dh.PixelPusher.GroupOrdinal = int32(ordinal.Group)
	}

	return dh
}

// DoneC implements device.D.
//
// A Proxy device is done when its base's DoneC has been closed.
func (pd *Device) DoneC() <-chan struct{} { return pd.doneC }

// Addr implements device.D.
func (pd *Device) Addr() net.Addr { return pd.proxy.Addr() }

// Info implements device.D.
func (pd *Device) Info() (i device.Info) {
	pd.modInfo(func(di *device.Info) {
		i = device.Info{
			PacketsReceived: di.PacketsReceived,
			BytesReceived:   di.BytesReceived,

			PacketsSent: di.PacketsSent,
			BytesSent:   di.BytesSent,

			Created:  pd.createdTime,
			Observed: pd.createdTime,
		}
	})
	return
}

// Proxied returns the base device that pd is proxying for.
func (pd *Device) Proxied() device.D { return pd.base }

// onLocalPacketData is called when pd.Local's callback receives data.
func (pd *Device) onLocalPacketData(buf *bufferpool.Buffer) {
	// Update our received metrics.
	pd.modInfo(func(di *device.Info) {
		di.PacketsReceived++
		di.BytesReceived += int64(buf.Len())
	})
	proxyReceivedPackets.With(pd.counterLabels).Inc()
	proxyReceivedBytes.With(pd.counterLabels).Add(float64(buf.Len()))

	// Dispatch the raw buffer to the proxied base, unless we're not forwarding.
	pkt := packetData{
		Buffer:    buf,
		forwarded: pd.m.Forwarding(),
	}
	if pkt.forwarded {
		buf.Retain()
		pd.basePacketsC <- &pkt
	} else {
		pd.logger.Debugf("NOT forwarding packet (size %d) to proxy device %v (forwarding is disabled).",
			buf.Len(), pd.baseID)
	}

	// Send the packet to our listeners channel.
	buf.Retain()
	pd.listenerPacketsC <- &pkt
}

// forwardPacketsToBase is run in its own goroutine. Its job is to forward raw
// packet data from basePacketsC to the base device immediately.
//
// forwardPacketsToBase takes ownership of pw, and will close it on exit.
func (pd *Device) forwardPacketsToBase(s device.Sender) {
	defer func() {
		if err := s.Close(); err != nil {
			pd.logger.Warnf("Failed to close %q Sender: %s", pd.baseID, err)
		}
	}()

	for pkt := range pd.basePacketsC {
		// Dispatch the packet to our underlying device.
		//
		// We do this in a separate scope so we can ensure that we release the
		// buffer on completion.
		func() {
			defer pkt.Release()

			if err := s.SendDatagram(pkt.Bytes()); err != nil {
				pd.logger.Warnf("Failed to forward packet from proxy device %q for: %s", pd.proxy.DeviceID, err)
				proxyForwardErrors.With(pd.counterLabels).Inc()
				return
			}

			// Update our sent metrics.
			pd.modInfo(func(di *device.Info) {
				di.PacketsSent++
				di.BytesSent += int64(pkt.Len())
			})
			proxySentPackets.With(pd.counterLabels).Inc()
			proxySentBytes.With(pd.counterLabels).Add(float64(pkt.Len()))
		}()
	}
}

// forwardPacketsToListeners is run in its own goroutine. Its job is to forward
// packets to registered listeners.
func (pd *Device) forwardPacketsToListeners(pr *protocol.PacketReader) {
	var parsed protocol.Packet
	for pkt := range pd.listenerPacketsC {
		// Do we have registered listeners? If not, ignore this packet.
		//
		// NOTE: There is a "race" here between this decision and actually sending
		// the packet; however, losing the race just means dropping a packet, so
		// this is fine.
		if !pd.m.hasListeners() {
			continue
		}

		// Parse the packet and send it to listeners.
		//
		// We do this in a separate scope so we can release the buffer on
		// completion.
		func() {
			defer pkt.Release()

			// Parse the Packet.
			//
			// NOTE the parsed packet may contain references to the underlying
			// buffer. The buffer must not be released until handling is finished.
			bsr := byteslicereader.R{Buffer: pkt.Bytes()}
			if err := pr.ReadPacket(&bsr, &parsed); err != nil {
				// Release ownership of the buffer, opening it up for reuse.
				pd.logger.Warnf("Discarding unknown packet for %q: %s", pd.proxy.DeviceID, err)
				pd.logger.Debugf("Discarded packet contents:\n%s", fmtutil.Hex(pkt.Bytes()))
				return
			}

			pd.m.sendPacketToListeners(pd.base, &parsed, pkt.forwarded)
		}()
	}
}

func (pd *Device) closeWhenBaseCloses() {
	// Wait for either this device's Context to close, or its base device to
	// close.
	select {
	case <-pd.doneC:
		// Shutdown externally.
	case <-pd.base.DoneC():
		// Our base device has closed.
		pd.shutdown()
	}
}

func (pd *Device) shutdown() {
	pd.shutdownOnce.Do(func() {
		// Close our local connection. This will stop packets from being forwarded
		// to our callback.
		if err := pd.proxy.Close(); err != nil {
			pd.logger.Warnf("Failed to close proxy device %q: %s", pd.proxy.DeviceID, err)
		}

		// At this point, our proxy has been closed, so all of our packet callbacks
		// have completed and we will no longer receive any additional callbacks.
		//
		// Shut down our processing goroutines.
		close(pd.basePacketsC)
		close(pd.listenerPacketsC)

		// Mark that we are done to external viewers.
		close(pd.doneC)

		pd.logger.Infof("Proxy device %q has finished; unregistering from manager.", pd.proxy.DeviceID)
		pd.m.removeDevice(pd)
	})

	// Wait for any outstanding processes to finish.
	pd.wg.Wait()
}

func (pd *Device) modInfo(fn func(*device.Info)) {
	pd.infoMu.Lock()
	defer pd.infoMu.Unlock()
	fn(&pd.info)
}

// packetData is data belonging to a packet.
type packetData struct {
	// buf if the buffer containing the packet.
	*bufferpool.Buffer
	// forwarded is true if the packet was forwarded to the underlying device.
	forwarded bool
}