ollie/cmd/olliesrv/main.go

// olliesrv - 9P server for ollie sessions
package main

import (
	"context"
	"flag"
	"fmt"
	"net"
	"os"
	"os/signal"
	"path/filepath"
	"syscall"

	"ollie/cmd/olliesrv/internal/agent"
	"ollie/cmd/olliesrv/internal/fs"
	olog "ollie/log"
	"ollie/util"

	"9fans.net/go/plan9/client"

	"github.com/godbus/dbus/v5"
)

const serviceName = "ollie"

var tcpAddr = flag.String("tcp", "", "also listen on TCP address (e.g. :564)")
var yolo = flag.Bool("yolo", false, "skip landrun sandbox for shell")

func main() {
	flag.CommandLine.Parse(os.Args[1:]) //nolint:errcheck

	ns := client.Namespace()
	if ns == "" {
		fmt.Fprintln(os.Stderr, "no namespace")
		os.Exit(1)
	}

	sockPath := filepath.Join(ns, serviceName)
	runServer(sockPath)
}

func runServer(sockPath string) {
	util.EnsureEnv()

	// Remove stale sockets from this and prior runs
	if _, err := os.Stat(sockPath); err == nil {
		os.Remove(sockPath) //nolint:errcheck
	}
	sockDir := filepath.Dir(sockPath)
	if entries, err := filepath.Glob(filepath.Join(sockDir, "toolsrv-*.sock")); err == nil {
		for _, s := range entries {
			os.Remove(s) //nolint:errcheck
		}
	}

	sink := olog.NewSink(os.Stdout, os.Stderr, olog.ParseLevel(os.Getenv("OLLIE_LOG"), olog.LevelWarn))

	agentsDirs := agent.AgentsDirs()
	sessionsDir := util.DataDir() + "/sessions"

	daemonCtx, daemonCancel := context.WithCancel(context.Background())
	defer daemonCancel()

	modelCache := fs.NewModelCache()

	// Desktop notifications for bypass prompts
	if conn, err := dbus.SessionBus(); err == nil {
		initBypassNotifier(conn)
	}

	rootTree := fs.NewRoot(fs.Config{
		Ctx:          daemonCtx,
		AgentsDir:    agentsDirs[0],
		SessionsDir:  sessionsDir,
		Log:          sink.NewLogger("9p"),
		Sink:         sink,
		Yolo:         *yolo,
		ModelCache:   modelCache,
		BypassNotify: notifyBypass,
	})

	// Create 9P server
	srv := New(Config{
		Sink:     sink,
		RootTree: rootTree,
	})

	// Prefetch models asynchronously at startup
	go modelCache.Get()

	var listener net.Listener
	var err error
	listener, err = net.Listen("unix", sockPath)
	if err != nil {
		fmt.Fprintf(os.Stderr, "listen: %v\n", err)
		os.Exit(1)
	}

	go func() {
		for {
			conn, err := listener.Accept()
			if err != nil {
				return
			}
			go srv.Start(conn)
		}
	}()

	fmt.Printf("olliesrv listening on %s\n", sockPath)

	// Restore sessions now that the 9P listener is ready.
	fs.RestoreSessions()

	var tcpListener net.Listener
	if *tcpAddr != "" {
		tcpListener, err = net.Listen("tcp", *tcpAddr)
		if err != nil {
			fmt.Fprintf(os.Stderr, "tcp listen: %v\n", err)
			os.Exit(1)
		}
		go func() {
			for {
				conn, err := tcpListener.Accept()
				if err != nil {
					return
				}
				go srv.Start(conn)
			}
		}()
		fmt.Printf("olliesrv listening on tcp %s\n", *tcpAddr)
	}

	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)
	<-sigChan

	fmt.Println("shutting down")

	// Stop accepting new connections first.
	if listener != nil {
		listener.Close() //nolint:errcheck
	}
	if tcpListener != nil {
		tcpListener.Close() //nolint:errcheck
	}

	daemonCancel() // signal all sessions via context propagation
	if srv != nil {
		srv.Kill()
	}
	os.Remove(sockPath)
	sink.Flush()
}