telebit/relay/api/conn_tracking.go

package api

import (
	"context"
	"fmt"
	"log"
	"net"
	"sync"
)

//Track -- used to track connection + domain
type Track struct {
	conn   net.Conn
	domain string
}

//NewTrack -- Constructor
func NewTrack(conn net.Conn, domain string) (p *Track) {
	p = new(Track)
	p.conn = conn
	p.domain = domain
	return
}

//Tracking --
type Tracking struct {
	mutex       sync.Mutex
	connections map[string]*Track
	register    chan *Track
	unregister  chan net.Conn
}

//NewTracking -- Constructor
func NewTracking() (p *Tracking) {
	p = new(Tracking)
	p.connections = make(map[string]*Track)
	p.register = make(chan *Track)
	p.unregister = make(chan net.Conn)
	return
}

//Run -
func (p *Tracking) Run(ctx context.Context) {
	log.Println("Tracking Running")

	for {
		select {

		case <-ctx.Done():
			log.Println("Cancel signal hit")
			return

		case connection := <-p.register:
			p.mutex.Lock()
			key := connection.conn.RemoteAddr().String()
			log.Println("register fired", key)
			p.connections[key] = connection
			p.list()
			p.mutex.Unlock()

		case connection := <-p.unregister:
			p.mutex.Lock()
			key := connection.RemoteAddr().String()
			log.Println("unregister fired", key)
			if _, ok := p.connections[key]; ok {
				delete(p.connections, key)
			}
			p.list()
			p.mutex.Unlock()
		}
	}
}

func (p *Tracking) list() {
	for c := range p.connections {
		log.Println(c)
	}
}

//Lookup --
// - get connection from key
func (p *Tracking) Lookup(key string) (*Track, error) {
	p.mutex.Lock()
	defer p.mutex.Unlock()

	if _, ok := p.connections[key]; ok {
		return p.connections[key], nil
	}
	return nil, fmt.Errorf("Lookup failed for %s", key)
}