repos / pico

pico services mono repo
git clone https://github.com/picosh/pico.git

pico / pkg / shared
Antonio Mika  ·  2025-03-12

senpai.go

  1package shared
  2
  3import (
  4	"fmt"
  5	"log/slog"
  6	"sync"
  7
  8	"git.sr.ht/~delthas/senpai"
  9	"github.com/containerd/console"
 10	"github.com/picosh/pico/pkg/pssh"
 11)
 12
 13type consoleData struct {
 14	data []byte
 15	err  error
 16}
 17
 18type VConsole struct {
 19	Session *pssh.SSHServerConnSession
 20	pty     *pssh.Pty
 21
 22	sizeEnableOnce sync.Once
 23
 24	windowMu      sync.Mutex
 25	currentWindow pssh.Window
 26
 27	readReq  chan []byte
 28	dataChan chan consoleData
 29
 30	wg sync.WaitGroup
 31}
 32
 33func (v *VConsole) Read(p []byte) (int, error) {
 34	v.wg.Add(1)
 35	defer v.wg.Done()
 36	tot := 0
 37
 38	if len(v.readReq) == 0 {
 39		select {
 40		case v.readReq <- make([]byte, len(p)):
 41		case <-v.Session.Context().Done():
 42			return 0, v.Session.Context().Err()
 43		}
 44	}
 45
 46	v.sizeEnableOnce.Do(func() {
 47		tot += copy(p, []byte("\x1b[?2048h"))
 48		v.windowMu.Lock()
 49		select {
 50		case v.dataChan <- consoleData{[]byte(fmt.Sprintf("\x1b[48;%d;%d;%d;%dt", v.currentWindow.Height, v.currentWindow.Width, v.currentWindow.HeightPixels, v.currentWindow.WidthPixels)), nil}:
 51		case <-v.Session.Context().Done():
 52			return
 53		}
 54		v.windowMu.Unlock()
 55		select {
 56		case v.dataChan <- consoleData{[]byte("\x1b[?1;0c"), nil}:
 57		case <-v.Session.Context().Done():
 58			return
 59		}
 60	})
 61
 62	if tot > 0 {
 63		return tot, nil
 64	}
 65
 66	select {
 67	case data := <-v.dataChan:
 68		tot += copy(p, data.data)
 69		return tot, data.err
 70	case <-v.Session.Context().Done():
 71		return 0, v.Session.Context().Err()
 72	}
 73}
 74
 75func (v *VConsole) Resize(winSize console.WinSize) error {
 76	return v.pty.Resize(int(winSize.Height), int(winSize.Width))
 77}
 78
 79func (v *VConsole) ResizeFrom(c console.Console) error {
 80	s, err := c.Size()
 81	if err != nil {
 82		return err
 83	}
 84
 85	return v.Resize(s)
 86}
 87
 88func (v *VConsole) SetRaw() error {
 89	return nil
 90}
 91
 92func (v *VConsole) DisableEcho() error {
 93	return nil
 94}
 95
 96func (v *VConsole) Reset() error {
 97	_, err := v.Session.Write([]byte("\033[?25h\033[0 q\033[34h\033[?25h\033[39;49m\033[m^O\033[H\033[J\033[?1049l\033[?1l\033>\033[?1000l\033[?1002l\033[?1003l\033[?1006l\033[?2004l"))
 98	return err
 99}
100
101func (v *VConsole) Size() (console.WinSize, error) {
102	v.windowMu.Lock()
103	defer v.windowMu.Unlock()
104	return console.WinSize{
105		Height: uint16(v.currentWindow.Height),
106		Width:  uint16(v.currentWindow.Width),
107	}, nil
108}
109
110func (v *VConsole) Fd() uintptr {
111	return 0
112}
113
114func (v *VConsole) Name() string {
115	return v.pty.Name()
116}
117
118func (v *VConsole) Close() error {
119	err := v.Session.Close()
120	v.wg.Wait()
121	close(v.readReq)
122	close(v.dataChan)
123	return err
124}
125
126func (v *VConsole) Write(p []byte) (int, error) {
127	return v.Session.Write(p)
128}
129
130func NewVConsole(sesh *pssh.SSHServerConnSession) (*VConsole, error) {
131	pty, win, ok := sesh.Pty()
132	if !ok {
133		return nil, fmt.Errorf("PTY not found")
134	}
135
136	vty := &VConsole{
137		pty:           pty,
138		Session:       sesh,
139		currentWindow: pty.Window,
140		readReq:       make(chan []byte, 100),
141		dataChan:      make(chan consoleData, 100),
142	}
143
144	vty.wg.Add(2)
145
146	go func() {
147		defer vty.wg.Done()
148		for {
149			select {
150			case <-sesh.Context().Done():
151				return
152			case data := <-vty.readReq:
153				n, err := vty.Session.Read(data)
154				select {
155				case vty.dataChan <- consoleData{data[:n], err}:
156				case <-sesh.Context().Done():
157					return
158				}
159			}
160		}
161	}()
162
163	go func() {
164		defer vty.wg.Done()
165		for {
166			select {
167			case <-sesh.Context().Done():
168				return
169			case w := <-win:
170				vty.windowMu.Lock()
171				vty.currentWindow = w
172				vty.windowMu.Unlock()
173				select {
174				case vty.dataChan <- consoleData{[]byte(fmt.Sprintf("\x1b[48;%d;%d;%d;%dt", w.Height, w.Width, w.HeightPixels, w.WidthPixels)), nil}:
175				case <-sesh.Context().Done():
176					return
177				}
178			}
179		}
180	}()
181
182	return vty, nil
183}
184
185func NewSenpaiApp(sesh *pssh.SSHServerConnSession, username, pass string) (*senpai.App, error) {
186	vty, err := NewVConsole(sesh)
187	if err != nil {
188		slog.Error("PTY not found")
189		return nil, err
190	}
191
192	senpaiCfg := senpai.Defaults()
193	senpaiCfg.TLS = true
194	senpaiCfg.Addr = "irc.pico.sh:6697"
195	senpaiCfg.Nick = username
196	senpaiCfg.Password = &pass
197	senpaiCfg.WithConsole = vty
198
199	app, err := senpai.NewApp(senpaiCfg)
200	return app, err
201}