package clients import ( "context" "crypto/rand" "encoding/json" "errors" "fmt" "log/slog" "slices" "sort" "strings" "sync" "time" "github.com/coder/websocket" "g.fc.run/misc/byre/internal/api" ) const ( readLimit = 256 << 20 pingInterval = 30 * time.Second pingTimeout = 15 * time.Second ) type Client struct { Login string Host string Commit string ConnectedAt time.Time } type Transfer struct { Dir string Name string Receive bool } type conn struct { Client ws *websocket.Conn done chan struct{} mu sync.Mutex pending map[string]chan api.ClientReply } type Registry struct { mu sync.Mutex conns map[string]*conn transfers map[string]Transfer } func New() *Registry { return &Registry{conns: map[string]*conn{}, transfers: map[string]Transfer{}} } func (r *Registry) Serve(ctx context.Context, ws *websocket.Conn, client Client) { ws.SetReadLimit(readLimit) c := &conn{Client: client, ws: ws, done: make(chan struct{}), pending: map[string]chan api.ClientReply{}} r.mu.Lock() old := r.conns[client.Host] r.conns[client.Host] = c r.mu.Unlock() if old != nil { old.ws.Close(websocket.StatusGoingAway, "replaced by a newer connection from this host") } slog.Info("client connected", "host", client.Host, "login", client.Login, "commit", client.Commit) defer func() { r.mu.Lock() if r.conns[client.Host] == c { delete(r.conns, client.Host) } r.mu.Unlock() close(c.done) slog.Info("client disconnected", "host", client.Host, "login", client.Login) }() ctx, cancel := context.WithCancel(ctx) defer cancel() go c.ping(ctx) for { _, b, err := ws.Read(ctx) if err != nil { return } reply := api.ClientReply{} err = json.Unmarshal(b, &reply) if err != nil { ws.Close(websocket.StatusUnsupportedData, "messages must be ClientReply") return } c.mu.Lock() ch := c.pending[reply.Id] delete(c.pending, reply.Id) c.mu.Unlock() if ch == nil { slog.Warn("client reply for no request", "host", client.Host, "id", reply.Id) continue } ch <- reply } } func (c *conn) ping(ctx context.Context) { t := time.NewTicker(pingInterval) defer t.Stop() for { select { case <-ctx.Done(): return case <-t.C: pctx, cancel := context.WithTimeout(ctx, pingTimeout) err := c.ws.Ping(pctx) cancel() if err != nil { c.ws.Close(websocket.StatusGoingAway, "ping timed out") return } } } } func (r *Registry) List() []Client { r.mu.Lock() defer r.mu.Unlock() out := []Client{} for _, c := range r.conns { out = append(out, c.Client) } sort.Slice(out, func(i, j int) bool { return out[i].Host < out[j].Host }) return out } func (r *Registry) Call(ctx context.Context, host string, req api.ClientRequest) (api.ClientReply, error) { r.mu.Lock() c := r.conns[host] hosts := []string{} for h := range r.conns { hosts = append(hosts, h) } r.mu.Unlock() if c == nil { slices.Sort(hosts) return api.ClientReply{}, fmt.Errorf("no client named %q is connected; connected: %s", host, strings.Join(hosts, ", ")) } req.Id = rand.Text() ch := make(chan api.ClientReply, 1) c.mu.Lock() c.pending[req.Id] = ch c.mu.Unlock() defer func() { c.mu.Lock() delete(c.pending, req.Id) c.mu.Unlock() }() data, err := json.Marshal(req) if err != nil { return api.ClientReply{}, err } err = c.ws.Write(ctx, websocket.MessageText, data) if err != nil { return api.ClientReply{}, fmt.Errorf("send to %s: %w", host, err) } select { case reply := <-ch: if reply.Error != nil { return reply, errors.New(*reply.Error) } return reply, nil case <-c.done: return api.ClientReply{}, fmt.Errorf("%s disconnected before replying", host) case <-ctx.Done(): return api.ClientReply{}, fmt.Errorf("%s didn't reply: %w", host, ctx.Err()) } } func (r *Registry) AddTransfer(t Transfer) string { id := rand.Text() r.mu.Lock() defer r.mu.Unlock() r.transfers[id] = t return id } func (r *Registry) RemoveTransfer(id string) { r.mu.Lock() defer r.mu.Unlock() delete(r.transfers, id) } func (r *Registry) Transfer(id string) (Transfer, bool) { r.mu.Lock() defer r.mu.Unlock() t, ok := r.transfers[id] return t, ok }