telebit/server/conn_tracking.go

90 lines
1.6 KiB
Go

package server
import (
"context"
"fmt"
"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) {
loginfo.Println("Tracking Running")
for {
select {
case <-ctx.Done():
loginfo.Println("Cancel signal hit")
return
case connection := <-p.register:
p.mutex.Lock()
key := connection.conn.RemoteAddr().String()
loginfo.Println("register fired", key)
p.connections[key] = connection
p.list()
p.mutex.Unlock()
case connection := <-p.unregister:
p.mutex.Lock()
key := connection.RemoteAddr().String()
loginfo.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 {
loginfo.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)
}