chore: include complete dc refactor implementation and audits
Release / Build & Release (push) Canceled after 0s

This commit is contained in:
kory
2026-10-02 13:23:30 +08:00
parent 192969d7f6
commit 38382d2486
22 changed files with 3022 additions and 8 deletions
+130
View File
@@ -0,0 +1,130 @@
package executor
import (
"crypto/sha256"
"encoding/hex"
"errors"
"fmt"
"os"
"path/filepath"
"sync"
)
var ErrBlockNotUnique = errors.New("old_text must occur exactly once")
var ErrEditConflict = errors.New("file changed during edit")
var editMu sync.Mutex
type EditBlockResult struct {
BytesWritten int64
}
func EditBlock(path, oldText, newText string, expectedHash ...string) (EditBlockResult, error) {
editMu.Lock()
defer editMu.Unlock()
if path == "" {
return EditBlockResult{}, errors.New("path required")
}
if oldText == "" {
return EditBlockResult{}, errors.New("old_text required")
}
data, err := os.ReadFile(path)
if err != nil {
return EditBlockResult{}, err
}
initial := sha256.Sum256(data)
initialHash := hex.EncodeToString(initial[:])
if len(expectedHash) > 0 && expectedHash[0] != "" && expectedHash[0] != initialHash {
return EditBlockResult{}, ErrEditConflict
}
src := string(data)
count := countNonOverlapping(src, oldText)
if count != 1 {
if count == 0 {
return EditBlockResult{}, fmt.Errorf("%w: found 0 matches", ErrBlockNotUnique)
}
return EditBlockResult{}, fmt.Errorf("%w: found %d matches", ErrBlockNotUnique, count)
}
updated := replaceOnce(src, oldText, newText)
info, err := os.Stat(path)
if err != nil {
return EditBlockResult{}, err
}
// Re-read immediately before commit. This is the CAS check against
// writers outside this process; editMu serializes our own writers.
latest, err := os.ReadFile(path)
if err != nil {
return EditBlockResult{}, err
}
latestSum := sha256.Sum256(latest)
if hex.EncodeToString(latestSum[:]) != initialHash {
return EditBlockResult{}, ErrEditConflict
}
dir := filepath.Dir(path)
tmp, err := os.CreateTemp(dir, ".vps-commander-edit-*")
if err != nil {
return EditBlockResult{}, err
}
tmpName := tmp.Name()
defer os.Remove(tmpName)
if err := tmp.Chmod(info.Mode().Perm()); err != nil {
_ = tmp.Close()
return EditBlockResult{}, err
}
if _, err := tmp.WriteString(updated); err != nil {
_ = tmp.Close()
return EditBlockResult{}, err
}
if err := tmp.Sync(); err != nil {
_ = tmp.Close()
return EditBlockResult{}, err
}
if err := tmp.Close(); err != nil {
return EditBlockResult{}, err
}
if err := os.Rename(tmpName, path); err != nil {
return EditBlockResult{}, err
}
if d, err := os.Open(dir); err == nil {
_ = d.Sync()
_ = d.Close()
}
return EditBlockResult{BytesWritten: int64(len(updated))}, nil
}
func countNonOverlapping(src, old string) int {
n, count := 0, 0
for {
i := indexString(src[n:], old)
if i < 0 {
return count
}
count++
n += i + len(old)
if n > len(src) {
return count
}
}
}
func replaceOnce(src, old, new string) string {
i := indexString(src, old)
return src[:i] + new + src[i+len(old):]
}
func indexString(src, needle string) int {
if len(needle) == 0 {
return 0
}
for i := 0; i+len(needle) <= len(src); i++ {
if src[i:i+len(needle)] == needle {
return i
}
}
return -1
}
+52
View File
@@ -0,0 +1,52 @@
package executor
import (
"errors"
"os"
"path/filepath"
"testing"
)
func TestEditBlockUniqueAndAtomic(t *testing.T) {
dir := t.TempDir()
path := filepath.Join(dir, "config.txt")
if err := os.WriteFile(path, []byte("a\nTARGET\nz\n"), 0640); err != nil {
t.Fatal(err)
}
res, err := EditBlock(path, "TARGET", "CHANGED")
if err != nil {
t.Fatal(err)
}
if res.BytesWritten != int64(len("a\nCHANGED\nz\n")) {
t.Fatalf("bytes=%d", res.BytesWritten)
}
got, _ := os.ReadFile(path)
if string(got) != "a\nCHANGED\nz\n" {
t.Fatalf("content=%q", got)
}
info, _ := os.Stat(path)
if info.Mode().Perm() != 0640 {
t.Fatalf("mode=%o, want 0640", info.Mode().Perm())
}
}
func TestEditBlockRejectsZeroAndMultipleMatchesWithoutMutation(t *testing.T) {
dir := t.TempDir()
path := filepath.Join(dir, "config.txt")
original := "TARGET\nother\nTARGET\n"
if err := os.WriteFile(path, []byte(original), 0644); err != nil {
t.Fatal(err)
}
_, err := EditBlock(path, "MISSING", "x")
if !errors.Is(err, ErrBlockNotUnique) {
t.Fatalf("missing match error=%v", err)
}
_, err = EditBlock(path, "TARGET", "x")
if !errors.Is(err, ErrBlockNotUnique) {
t.Fatalf("multiple match error=%v", err)
}
got, _ := os.ReadFile(path)
if string(got) != original {
t.Fatalf("file mutated after rejected edit: %q", got)
}
}
+36
View File
@@ -0,0 +1,36 @@
package executor
import (
"context"
"os"
"path/filepath"
"testing"
)
func TestListDirectoryAndInfo(t *testing.T) {
d := t.TempDir()
if err := os.WriteFile(filepath.Join(d, "a.txt"), []byte("hello"), 0600); err != nil {
t.Fatal(err)
}
if err := os.Mkdir(filepath.Join(d, "sub"), 0700); err != nil {
t.Fatal(err)
}
es, err := ListDirectory(d, false, 10)
if err != nil || len(es) != 2 {
t.Fatalf("entries=%v err=%v", es, err)
}
fi, err := GetFileInfo(filepath.Join(d, "a.txt"))
if err != nil || fi.Size != 5 {
t.Fatalf("info=%+v err=%v", fi, err)
}
}
func TestSearchHonorsContextAndLimit(t *testing.T) {
d := t.TempDir()
for i := 0; i < 3; i++ {
_ = os.WriteFile(filepath.Join(d, string(rune('a'+i))+".txt"), []byte("needle\n"), 0600)
}
r, err := Search(context.Background(), d, "needle", 2)
if err != nil || len(r) != 2 {
t.Fatalf("results=%v err=%v", r, err)
}
}
+248
View File
@@ -0,0 +1,248 @@
package executor
import (
"context"
"errors"
"io"
"os/exec"
"sync"
"syscall"
"time"
)
const (
DefaultRingCapacity = 2 * 1024 * 1024
MaxSessionInput = 64 * 1024
DefaultSessionTTL = 30 * time.Minute
SessionIdleTTL = 10 * time.Minute
MaxSessionBufferTotal = 8 * 1024 * 1024
)
type RingBuffer struct {
mu sync.RWMutex
buf []byte
head int
size int
}
func NewRingBuffer(capacity int) *RingBuffer {
if capacity <= 0 {
capacity = DefaultRingCapacity
}
return &RingBuffer{buf: make([]byte, capacity)}
}
func (r *RingBuffer) Write(p []byte) (int, error) {
total := len(p)
r.mu.Lock()
defer r.mu.Unlock()
for len(p) > 0 {
n := len(p)
if n > len(r.buf) {
p = p[n-len(r.buf):]
n = len(p)
}
first := len(r.buf) - r.head
if first > n {
first = n
}
copy(r.buf[r.head:r.head+first], p[:first])
if first < n {
copy(r.buf[:n-first], p[first:n])
}
r.head = (r.head + n) % len(r.buf)
r.size += n
if r.size > len(r.buf) {
r.size = len(r.buf)
}
p = p[n:]
}
return total, nil
}
func (r *RingBuffer) Snapshot() []byte {
r.mu.RLock()
defer r.mu.RUnlock()
out := make([]byte, r.size)
start := (r.head - r.size + len(r.buf)) % len(r.buf)
first := len(r.buf) - start
if first > r.size {
first = r.size
}
copy(out, r.buf[start:start+first])
if first < r.size {
copy(out[first:], r.buf[:r.size-first])
}
return out
}
type Session struct {
ID string
Command string
Owner string
StartedAt time.Time
ExpiresAt time.Time
LastAccess time.Time
cmd *exec.Cmd
stdin io.WriteCloser
stdout *RingBuffer
stderr *RingBuffer
done chan struct{}
mu sync.RWMutex
err error
exitCode int
}
func StartSession(ctx context.Context, id, command, workdir, owner string, capacity int) (*Session, error) {
if id == "" || command == "" {
return nil, errors.New("session id and command required")
}
cmd := exec.Command("/bin/sh", "-c", command)
cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true}
if workdir != "" {
cmd.Dir = workdir
}
stdout := NewRingBuffer(capacity)
stderr := NewRingBuffer(capacity)
stdin, err := cmd.StdinPipe()
if err != nil {
return nil, err
}
out, err := cmd.StdoutPipe()
if err != nil {
return nil, err
}
errOut, err := cmd.StderrPipe()
if err != nil {
return nil, err
}
now := time.Now()
s := &Session{
ID: id, Command: command, Owner: owner, StartedAt: now, LastAccess: now, ExpiresAt: now.Add(DefaultSessionTTL), cmd: cmd,
stdin: stdin, stdout: stdout, stderr: stderr, done: make(chan struct{}), exitCode: -1,
}
if err := cmd.Start(); err != nil {
_ = stdin.Close()
return nil, err
}
go s.copy(out, stdout)
go s.copy(errOut, stderr)
go s.wait()
go func() {
select {
case <-ctx.Done():
s.Terminate()
case <-s.done:
}
}()
return s, nil
}
func (s *Session) copy(src io.Reader, dst io.Writer) {
_, _ = io.Copy(dst, src)
}
func (s *Session) wait() {
err := s.cmd.Wait()
s.mu.Lock()
s.err = err
if err == nil {
s.exitCode = 0
} else if ee, ok := err.(*exec.ExitError); ok {
s.exitCode = ee.ExitCode()
} else {
s.exitCode = 1
}
s.mu.Unlock()
close(s.done)
}
func (s *Session) Done() <-chan struct{} { return s.done }
func (s *Session) Status() (int, bool, error) {
select {
case <-s.done:
s.mu.RLock()
defer s.mu.RUnlock()
return s.exitCode, true, s.err
default:
return -1, false, nil
}
}
func (s *Session) Output() (stdout, stderr string) {
stdout, stderr = string(s.stdout.Snapshot()), string(s.stderr.Snapshot())
s.mu.Lock()
if time.Now().Before(s.ExpiresAt) {
s.LastAccess = time.Now()
}
s.mu.Unlock()
return stdout, stderr
}
func (s *Session) Stdin(data []byte) error {
if len(data) > MaxSessionInput {
return errors.New("session input too large")
}
s.mu.RLock()
stdin := s.stdin
s.mu.RUnlock()
if stdin == nil {
return errors.New("session stdin unavailable")
}
_, err := stdin.Write(data)
if err == nil {
s.mu.Lock()
s.LastAccess = time.Now()
s.mu.Unlock()
}
return err
}
func (s *Session) Authorized(owner string) bool {
s.mu.Lock()
defer s.mu.Unlock()
now := time.Now()
if owner == "" || s.Owner != owner || now.After(s.ExpiresAt) || now.Sub(s.LastAccess) > SessionIdleTTL {
return false
}
s.LastAccess = now
return true
}
func (s *Session) Expired() bool {
s.mu.RLock()
defer s.mu.RUnlock()
now := time.Now()
return now.After(s.ExpiresAt) || now.Sub(s.LastAccess) > SessionIdleTTL
}
func (s *Session) Terminate() error {
s.mu.RLock()
pid := 0
if s.cmd != nil && s.cmd.Process != nil {
pid = s.cmd.Process.Pid
}
s.mu.RUnlock()
if pid == 0 {
return nil
}
return syscall.Kill(-pid, syscall.SIGTERM)
}
func (s *Session) Kill() error {
s.mu.RLock()
pid := 0
if s.cmd != nil && s.cmd.Process != nil {
pid = s.cmd.Process.Pid
}
s.mu.RUnlock()
if pid == 0 {
return nil
}
return syscall.Kill(-pid, syscall.SIGKILL)
}
// Compile-time sanity for the output writer contract.
var _ io.Writer = (*RingBuffer)(nil)
+104
View File
@@ -0,0 +1,104 @@
package executor
import (
"context"
"errors"
"sync"
"time"
)
const MaxSessionsPerManager = MaxSessionBufferTotal / DefaultRingCapacity
type SessionInfo struct {
ID string `json:"id"`
Command string `json:"command"`
StartedAt time.Time `json:"started_at"`
Running bool `json:"running"`
ExitCode int `json:"exit_code"`
}
type SessionManager struct {
mu sync.RWMutex
sessions map[string]*Session
limit int
}
func NewSessionManager(limit int) *SessionManager {
if limit <= 0 || limit > MaxSessionsPerManager {
limit = MaxSessionsPerManager
}
return &SessionManager{sessions: make(map[string]*Session), limit: limit}
}
func (m *SessionManager) Start(ctx context.Context, id, command, workdir, owner string, capacity int) (*Session, error) {
m.mu.Lock()
if _, ok := m.sessions[id]; ok {
m.mu.Unlock()
return nil, errors.New("session already exists")
}
if len(m.sessions) >= m.limit {
m.mu.Unlock()
return nil, errors.New("session limit reached")
}
s, err := StartSession(ctx, id, command, workdir, owner, capacity)
if err != nil {
m.mu.Unlock()
return nil, err
}
if _, exists := m.sessions[id]; exists {
m.mu.Unlock()
_ = s.Kill()
return nil, errors.New("session already exists")
}
m.sessions[id] = s
m.mu.Unlock()
return s, nil
}
func (m *SessionManager) Get(id string) (*Session, bool) {
m.mu.RLock()
defer m.mu.RUnlock()
s, ok := m.sessions[id]
return s, ok
}
func (m *SessionManager) Remove(id string) bool {
m.mu.Lock()
defer m.mu.Unlock()
if _, ok := m.sessions[id]; !ok {
return false
}
delete(m.sessions, id)
return true
}
func (m *SessionManager) List() []SessionInfo {
m.mu.RLock()
defer m.mu.RUnlock()
out := make([]SessionInfo, 0, len(m.sessions))
for _, s := range m.sessions {
code, done, _ := s.Status()
out = append(out, SessionInfo{ID: s.ID, Command: s.Command, StartedAt: s.StartedAt, Running: !done, ExitCode: code})
}
return out
}
func (m *SessionManager) Stop(id string, force bool) error {
s, ok := m.Get(id)
if !ok {
return errors.New("session not found")
}
if force {
return s.Kill()
}
return s.Terminate()
}
func (m *SessionManager) Cleanup() {
m.mu.Lock()
defer m.mu.Unlock()
for id, s := range m.sessions {
if s.Expired() {
_, done, _ := s.Status()
if !done {
_ = s.Kill()
continue
}
delete(m.sessions, id)
}
}
}
+72
View File
@@ -0,0 +1,72 @@
package executor
import (
"context"
"strings"
"testing"
"time"
)
func TestRingBufferKeepsNewestData(t *testing.T) {
r := NewRingBuffer(8)
_, _ = r.Write([]byte("12345678"))
_, _ = r.Write([]byte("90"))
if got := string(r.Snapshot()); got != "34567890" {
t.Fatalf("got %q", got)
}
}
func TestSessionCapturesAndTerminates(t *testing.T) {
s, err := StartSession(context.Background(), "t1", "printf hello; sleep 5", "", "owner", 1024)
if err != nil {
t.Fatal(err)
}
time.Sleep(100 * time.Millisecond)
out, _ := s.Output()
if !strings.Contains(out, "hello") {
t.Fatalf("out=%q", out)
}
if err := s.Terminate(); err != nil {
t.Fatal(err)
}
select {
case <-s.Done():
case <-time.After(2 * time.Second):
t.Fatal("session did not terminate")
}
}
func TestSessionManagerDuplicateAndLimit(t *testing.T) {
m := NewSessionManager(1)
s, err := m.Start(context.Background(), "a", "sleep 1", "", "owner", 128)
if err != nil {
t.Fatal(err)
}
defer s.Kill()
if _, err := m.Start(context.Background(), "b", "true", "", "owner", 128); err == nil {
t.Fatal("expected limit")
}
if _, err := m.Start(context.Background(), "a", "true", "", "owner", 128); err == nil {
t.Fatal("expected duplicate")
}
}
func TestSessionIdleAuthorizationTouch(t *testing.T) {
s, err := StartSession(context.Background(), "idle-test", "sleep 2", "", "owner", 128)
if err != nil {
t.Fatal(err)
}
defer s.Kill()
s.mu.Lock()
s.LastAccess = time.Now().Add(-(SessionIdleTTL + time.Second))
s.mu.Unlock()
if s.Authorized("owner") {
t.Fatal("idle-expired session must not authorize")
}
s.mu.Lock()
s.LastAccess = time.Now()
s.mu.Unlock()
if !s.Authorized("owner") {
t.Fatal("active session should authorize")
}
}