go-tunnel/cmd/server/registry.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(),
	}
}