package agentsocket

import (

"context"

"errors"

"net"

"sync"

"golang.org/x/xerrors"

"storj.io/drpc/drpcmux"

"storj.io/drpc/drpcserver"

"cdr.dev/slog/v3"

"github.com/coder/coder/v2/agent/agentsocket/proto"

agentproto "github.com/coder/coder/v2/agent/proto"

"github.com/coder/coder/v2/agent/unit"

"github.com/coder/coder/v2/codersdk/drpcsdk"

)

// Server provides access to the DRPCAgentSocketService via a Unix domain socket.

// Do not invoke Server{} directly. Use NewServer() instead.

type Server struct {

logger slog.Logger

path string

drpcServer *drpcserver.Server

service *DRPCAgentSocketService

mu sync.Mutex

listener net.Listener

ctx context.Context

cancel context.CancelFunc

wg sync.WaitGroup

}

// NewServer creates a new agent socket server.

func NewServer(logger slog.Logger, opts ...Option) (*Server, error) {

options := &options{}

for _, opt := range opts {

opt(options)

}

logger = logger.Named("agentsocket-server")

server := &Server{

logger: logger,

path: options.path,

service: &DRPCAgentSocketService{

logger: logger,

unitManager: unit.NewManager(),

contextManager: options.contextManager,

},

}

mux := drpcmux.New()

err := proto.DRPCRegisterAgentSocket(mux, server.service)

if err != nil {

return nil, xerrors.Errorf("failed to register drpc service: %w", err)

}

server.drpcServer = drpcserver.NewWithOptions(mux, drpcserver.Options{

Manager: drpcsdk.DefaultDRPCOptions(nil),

Log: func(err error) {

if errors.Is(err, context.Canceled) ||

errors.Is(err, context.DeadlineExceeded) {

return

}

logger.Debug(context.Background(), "drpc server error", slog.Error(err))

},

})

listener, err := createSocket(server.path)

if err != nil {

return nil, xerrors.Errorf("create socket: %w", err)

}

server.listener = listener

// This context is canceled by server.Close().

// canceling it will close all connections.

server.ctx, server.cancel = context.WithCancel(context.Background())

server.logger.Info(server.ctx, "agent socket server started", slog.F("path", server.path))

server.wg.Add(1)

go func() {

defer server.wg.Done()

server.acceptConnections()

}()

return server, nil

}

// Close stops the server and cleans up resources.

func (s *Server) Close() error {

s.mu.Lock()

if s.listener == nil {

s.mu.Unlock()

return nil

}

s.logger.Info(s.ctx, "stopping agent socket server")

s.cancel()

if err := s.listener.Close(); err != nil {

s.logger.Warn(s.ctx, "error closing socket listener", slog.Error(err))

}

s.listener = nil

s.mu.Unlock()

// Wait for all connections to finish

s.wg.Wait()

if err := cleanupSocket(s.path); err != nil {

s.logger.Warn(s.ctx, "error cleaning up socket file", slog.Error(err))

}

s.logger.Info(s.ctx, "agent socket server stopped")

return nil

}

// SetAgentAPI sets the agent API client used to forward requests

// to coderd.

func (s *Server) SetAgentAPI(api agentproto.DRPCAgentClient28) {

s.service.SetAgentAPI(api)

}

// ClearAgentAPI clears the agent API client.

func (s *Server) ClearAgentAPI() {

s.service.ClearAgentAPI()

}

func (s *Server) acceptConnections() {

// In an edge case, Close() might race with acceptConnections() and set s.listener to nil.

// Therefore, we grab a copy of the listener under a lock. We might still get a nil listener,

// but then we know close has already run and we can return early.

s.mu.Lock()

listener := s.listener

s.mu.Unlock()

if listener == nil {

return

}

err := s.drpcServer.Serve(s.ctx, listener)

if err != nil {

s.logger.Warn(s.ctx, "error serving drpc server", slog.Error(err))

}

}