gotermix/internals/session.go

package internals

import (
	"crypto/rand"
	"fmt"
	"os/exec"
	"runtime"
	"strings"
	"sync"
	"time"

	"github.com/creack/pty"
	"github.com/gorilla/websocket"
)

var (
	sessions   = map[string]*Session{}
	sessionsMu sync.Mutex
)

func shellQuote(s string) string {
	return "'" + strings.ReplaceAll(s, "'", `'\''`) + "'"
}

func validID(id string) bool {
	if len(id) != 32 {
		return false
	}
	for _, c := range id {
		if !((c >= '0' && c <= '9') || (c >= 'a' && c <= 'f')) {
			return false
		}
	}
	return true
}

func randHex(n int) string {
	b := make([]byte, n)
	rand.Read(b)
	return fmt.Sprintf("%x", b)
}

func getOrCreate(id string) *Session {
	sessionsMu.Lock()
	defer sessionsMu.Unlock()

	if s, ok := sessions[id]; ok {
		select {
		case <-s.done:
			delete(sessions, id)
		default:
			return s
		}
	}

	s := &Session{
		id:       id,
		clients:  make(map[*client]struct{}),
		done:     make(chan struct{}),
		lastSeen: time.Now(),
	}

	var cmd *exec.Cmd
	if runtime.GOOS == "windows" {
		cmd = exec.Command("cmd.exe")
	} else {
		cmd = exec.Command("/bin/bash", "-i")
	}
	cmd.Dir = initialCwd

	ptty, err := pty.Start(cmd)
	if err != nil {
		return nil
	}
	s.ptty = ptty
	s.cmd = cmd
	sessions[id] = s
	go s.readLoop()
	return s
}

func sessionByID(id string) *Session {
	sessionsMu.Lock()
	defer sessionsMu.Unlock()
	s, ok := sessions[id]
	if !ok {
		return nil
	}
	select {
	case <-s.done:
		return nil
	default:
		return s
	}
}

func (s *Session) readLoop() {
	chunk := make([]byte, 16*1024)
	for {
		n, err := s.ptty.Read(chunk)
		if n > 0 {
			data := make([]byte, n)
			copy(data, chunk[:n])

			s.mu.Lock()
			s.buf = append(s.buf, data...)
			if len(s.buf) > maxBufSize {
				s.buf = s.buf[len(s.buf)-maxBufSize:]
			}
			snapshot := make([]*client, 0, len(s.clients))
			for c := range s.clients {
				snapshot = append(snapshot, c)
			}
			s.mu.Unlock()

			for _, c := range snapshot {
				c.write(websocket.BinaryMessage, data)
			}
		}
		if err != nil {
			break
		}
	}
	close(s.done)
	sessionsMu.Lock()
	delete(sessions, s.id)
	sessionsMu.Unlock()
}