corago/codex-keep/rpc.go

package main

import (
	"context"
	"encoding/json"
	"fmt"
	"net"
	"net/http"
	"strconv"
	"time"

	"github.com/coder/websocket"
)

type frame struct {
	message message
	err     error
}

type connection struct {
	write   func([]byte) error
	frames  chan frame
	serial  uint64
	onEvent func(message) error
	stop    func()
}

func connect(ctx context.Context, socket string) (*connection, error) {
	ctx, cancel := context.WithCancel(ctx)
	transport := &http.Transport{DialContext: func(ctx context.Context, _, _ string) (net.Conn, error) {
		return (&net.Dialer{}).DialContext(ctx, "unix", socket)
	}}
	handshake, stop := context.WithTimeout(ctx, 5*time.Second)
	defer stop()
	ws, _, err := websocket.Dial(handshake, "ws://localhost/", &websocket.DialOptions{HTTPClient: &http.Client{Transport: transport}})
	if err != nil {
		cancel()
		transport.CloseIdleConnections()
		return nil, err
	}
	ws.SetReadLimit(-1)
	c := &connection{frames: make(chan frame, 16)}
	c.write = func(data []byte) error { return ws.Write(ctx, websocket.MessageText, data) }
	done := make(chan struct{})
	go func() {
		defer close(done)
		for {
			_, data, err := ws.Read(ctx)
			var m message
			if err == nil {
				err = json.Unmarshal(data, &m)
			}
			select {
			case c.frames <- frame{m, err}:
			case <-ctx.Done():
				return
			}
			if err != nil {
				return
			}
		}
	}()
	c.stop = func() {
		cancel()
		_ = ws.CloseNow()
		transport.CloseIdleConnections()
		<-done
	}
	return c, nil
}

func (c *connection) send(id uint64, method string, params any) error {
	data, err := json.Marshal(struct {
		ID     uint64 `json:"id,omitempty"`
		Method string `json:"method"`
		Params any    `json:"params"`
	}{id, method, params})
	if err != nil {
		return err
	}
	return c.write(data)
}

func (c *connection) event(m message) error {
	// Approval and input requests belong to the interactive clients.
	if len(m.ID) != 0 || m.Method == "" || c.onEvent == nil {
		return nil
	}
	return c.onEvent(m)
}

func (c *connection) call(ctx context.Context, method string, params, result any) error {
	ctx, cancel := context.WithTimeout(ctx, 30*time.Second)
	defer cancel()
	c.serial++
	id := c.serial
	if err := c.send(id, method, params); err != nil {
		return err
	}
	for {
		select {
		case <-ctx.Done():
			return ctx.Err()
		case f := <-c.frames:
			if f.err != nil {
				return f.err
			}
			m := f.message
			if m.Method == "" && string(m.ID) == strconv.FormatUint(id, 10) {
				if m.Error != nil {
					return m.Error
				}
				if result == nil {
					return nil
				}
				if err := json.Unmarshal(m.Result, result); err != nil {
					return fmt.Errorf("%s: %w", method, err)
				}
				return nil
			}
			if err := c.event(m); err != nil {
				return err
			}
		}
	}
}