go-tunnel/cmd/server/registry.go

78 lines
1.6 KiB
Go
Raw Normal View History

2020-10-21 18:00:15 +02:00
package main
import (
"context"
"errors"
"forge.cadoles.com/wpetit/go-tunnel"
cmap "github.com/streamrail/concurrent-map"
)
type Registry struct {
clientIDByRemoteAddr cmap.ConcurrentMap
clientByClientID cmap.ConcurrentMap
}
func (r *Registry) OnClientAuth(ctx context.Context, rc *tunnel.RemoteClient, credentials interface{}) (bool, error) {
clientID, ok := credentials.(string)
if !ok {
return false, errors.New("unexpected credentials")
}
remoteAddr := rc.RemoteAddr().String()
r.Add(clientID, remoteAddr, rc)
return true, nil
}
func (r *Registry) OnClientDisconnect(ctx context.Context, rc *tunnel.RemoteClient) error {
remoteAddr := rc.RemoteAddr().String()
r.RemoveByRemoteAddr(remoteAddr)
return nil
}
func (r *Registry) Get(clientID string) *tunnel.RemoteClient {
rawRemoteClient, exists := r.clientByClientID.Get(clientID)
if !exists {
return nil
}
remoteClient, ok := rawRemoteClient.(*tunnel.RemoteClient)
if !ok {
return nil
}
return remoteClient
}
func (r *Registry) Add(clientID, remoteAddr string, rc *tunnel.RemoteClient) {
r.clientByClientID.Set(clientID, rc)
r.clientIDByRemoteAddr.Set(remoteAddr, clientID)
}
func (r *Registry) RemoveByRemoteAddr(remoteAddr string) {
rawClientID, exists := r.clientIDByRemoteAddr.Get(remoteAddr)
if !exists {
return
}
r.clientIDByRemoteAddr.Remove(remoteAddr)
clientID, ok := rawClientID.(string)
if !ok {
return
}
r.clientByClientID.Remove(clientID)
}
func NewRegistry() *Registry {
return &Registry{
clientIDByRemoteAddr: cmap.New(),
clientByClientID: cmap.New(),
}
}