78 lines
1.6 KiB
Go
78 lines
1.6 KiB
Go
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(),
|
|
}
|
|
}
|