new upload system

2026-03-08 00:53:22 -06:00
parent 8662f4310f
commit 4b48f462da
8 changed files with 1075 additions and 285 deletions
@@ -1345,6 +1345,12 @@ func HandleCommand(ctx context.Context, env *runtime.Env, envelope map[string]in
 			data = d
 		}
 		return HandleFileUpload(ctx, env, cmdID, path, data, offset, total, transferID)
+	case "file_upload_http":
+		payload, _ := envelope["payload"].(map[string]interface{})
+		path, _ := payload["path"].(string)
+		sourceURL, _ := payload["url"].(string)
+		total := payloadNumberToInt64(payload["total"])
+		return HandleFileUploadHTTP(ctx, env, cmdID, path, sourceURL, total)
 	case "file_delete":
 		path, _ := envelopePayloadString(envelope, "path")
 		return HandleFileDelete(ctx, env, cmdID, path)
@@ -5,9 +5,13 @@ import (
 	"bufio"
 	"bytes"
 	"context"
+	"crypto/tls"
+	"crypto/x509"
 	"fmt"
 	"io"
 	"log"
+	"net/http"
+	"net/url"
 	"os"
 	"path/filepath"
 	"runtime"
@@ -462,6 +466,88 @@ func HandleFileUpload(ctx context.Context, env *agentRuntime.Env, cmdID string,
 	return wire.WriteMsg(ctx, env.Conn, result)
 }

+func HandleFileUploadHTTP(ctx context.Context, env *agentRuntime.Env, cmdID string, destPath string, sourceURL string, expectedSize int64) error {
+	parsed, err := url.Parse(strings.TrimSpace(sourceURL))
+	if err != nil || parsed == nil || parsed.Host == "" {
+		return wire.WriteMsg(ctx, env.Conn, wire.CommandResult{Type: "command_result", CommandID: cmdID, OK: false, Message: "invalid upload url"})
+	}
+	if parsed.Scheme != "https" && parsed.Scheme != "http" {
+		return wire.WriteMsg(ctx, env.Conn, wire.CommandResult{Type: "command_result", CommandID: cmdID, OK: false, Message: "unsupported upload url scheme"})
+	}
+
+	tlsConfig := &tls.Config{InsecureSkipVerify: env.Cfg.TLSInsecureSkipVerify, MinVersion: tls.VersionTLS12}
+	if caPath := strings.TrimSpace(env.Cfg.TLSCAPath); caPath != "" {
+		caBytes, err := os.ReadFile(caPath)
+		if err != nil {
+			return wire.WriteMsg(ctx, env.Conn, wire.CommandResult{Type: "command_result", CommandID: cmdID, OK: false, Message: fmt.Sprintf("failed to read TLS CA: %v", err)})
+		}
+		pool := x509.NewCertPool()
+		if !pool.AppendCertsFromPEM(caBytes) {
+			return wire.WriteMsg(ctx, env.Conn, wire.CommandResult{Type: "command_result", CommandID: cmdID, OK: false, Message: "failed to parse TLS CA"})
+		}
+		tlsConfig.RootCAs = pool
+	}
+
+	client := &http.Client{Transport: &http.Transport{TLSClientConfig: tlsConfig}}
+	req, err := http.NewRequestWithContext(ctx, http.MethodGet, parsed.String(), nil)
+	if err != nil {
+		return wire.WriteMsg(ctx, env.Conn, wire.CommandResult{Type: "command_result", CommandID: cmdID, OK: false, Message: err.Error()})
+	}
+	if token := strings.TrimSpace(env.Cfg.AgentToken); token != "" {
+		req.Header.Set("x-agent-token", token)
+	}
+	if id := strings.TrimSpace(env.Cfg.ID); id != "" {
+		req.Header.Set("x-overlord-client-id", id)
+	}
+
+	resp, err := client.Do(req)
+	if err != nil {
+		return wire.WriteMsg(ctx, env.Conn, wire.CommandResult{Type: "command_result", CommandID: cmdID, OK: false, Message: err.Error()})
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
+		return wire.WriteMsg(ctx, env.Conn, wire.CommandResult{Type: "command_result", CommandID: cmdID, OK: false, Message: fmt.Sprintf("upload fetch failed: status %d", resp.StatusCode)})
+	}
+
+	dir := filepath.Dir(destPath)
+	if dir != "." {
+		if err := os.MkdirAll(dir, 0755); err != nil {
+			return wire.WriteMsg(ctx, env.Conn, wire.CommandResult{Type: "command_result", CommandID: cmdID, OK: false, Message: err.Error()})
+		}
+	}
+
+	tmpPath := destPath + ".httpuploading"
+	f, err := os.OpenFile(tmpPath, os.O_CREATE|os.O_TRUNC|os.O_WRONLY, 0644)
+	if err != nil {
+		return wire.WriteMsg(ctx, env.Conn, wire.CommandResult{Type: "command_result", CommandID: cmdID, OK: false, Message: err.Error()})
+	}
+
+	written, copyErr := io.Copy(f, resp.Body)
+	closeErr := f.Close()
+	if copyErr != nil {
+		_ = os.Remove(tmpPath)
+		return wire.WriteMsg(ctx, env.Conn, wire.CommandResult{Type: "command_result", CommandID: cmdID, OK: false, Message: copyErr.Error()})
+	}
+	if closeErr != nil {
+		_ = os.Remove(tmpPath)
+		return wire.WriteMsg(ctx, env.Conn, wire.CommandResult{Type: "command_result", CommandID: cmdID, OK: false, Message: closeErr.Error()})
+	}
+
+	if expectedSize > 0 && written != expectedSize {
+		_ = os.Remove(tmpPath)
+		return wire.WriteMsg(ctx, env.Conn, wire.CommandResult{Type: "command_result", CommandID: cmdID, OK: false, Message: "upload size mismatch"})
+	}
+
+	_ = os.Remove(destPath)
+	if err := os.Rename(tmpPath, destPath); err != nil {
+		_ = os.Remove(tmpPath)
+		return wire.WriteMsg(ctx, env.Conn, wire.CommandResult{Type: "command_result", CommandID: cmdID, OK: false, Message: err.Error()})
+	}
+
+	return wire.WriteMsg(ctx, env.Conn, wire.CommandResult{Type: "command_result", CommandID: cmdID, OK: true})
+}
+
 func HandleFileDelete(ctx context.Context, env *agentRuntime.Env, cmdID string, path string) error {
 	log.Printf("file_delete: %s", path)

@@ -3,6 +3,8 @@ package handlers
 import (
 	"bytes"
 	"context"
+	"net/http"
+	"net/http/httptest"
 	"os"
 	"path/filepath"
 	"testing"
@@ -467,3 +469,87 @@ func TestHandleFileZip(t *testing.T) {
 		t.Error("Expected either command_result or file_download message")
 	}
 }
+
+func TestHandleFileUploadHTTP(t *testing.T) {
+	data := []byte("http-upload-payload")
+
+	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if got := r.Header.Get("x-overlord-client-id"); got != "client-abc" {
+			t.Fatalf("expected x-overlord-client-id client-abc, got %q", got)
+		}
+		if got := r.Header.Get("x-agent-token"); got != "agent-token" {
+			t.Fatalf("expected x-agent-token agent-token, got %q", got)
+		}
+		_, _ = w.Write(data)
+	}))
+	defer ts.Close()
+
+	tmpDir := t.TempDir()
+	destPath := filepath.Join(tmpDir, "uploaded-http.bin")
+
+	writer := &testWriter{}
+	env := &rt.Env{
+		Conn: writer,
+		Cfg: config.Config{
+			ID:                    "client-abc",
+			AgentToken:            "agent-token",
+			TLSInsecureSkipVerify: true,
+		},
+	}
+
+	if err := HandleFileUploadHTTP(context.Background(), env, "cmd-http-upload", destPath, ts.URL+"/file", int64(len(data))); err != nil {
+		t.Fatalf("HandleFileUploadHTTP failed: %v", err)
+	}
+
+	content, err := os.ReadFile(destPath)
+	if err != nil {
+		t.Fatalf("failed reading uploaded file: %v", err)
+	}
+	if !bytes.Equal(content, data) {
+		t.Fatalf("uploaded content mismatch")
+	}
+
+	if len(writer.msgs) != 1 {
+		t.Fatalf("expected 1 message, got %d", len(writer.msgs))
+	}
+	var result wire.CommandResult
+	if err := msgpack.Unmarshal(writer.msgs[0], &result); err != nil {
+		t.Fatalf("failed to unmarshal command result: %v", err)
+	}
+	if !result.OK {
+		t.Fatalf("expected OK=true, got false: %s", result.Message)
+	}
+}
+
+func TestHandleFileUploadHTTP_SizeMismatch(t *testing.T) {
+	data := []byte("small")
+
+	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		_, _ = w.Write(data)
+	}))
+	defer ts.Close()
+
+	tmpDir := t.TempDir()
+	destPath := filepath.Join(tmpDir, "uploaded-http-mismatch.bin")
+
+	writer := &testWriter{}
+	env := &rt.Env{
+		Conn: writer,
+		Cfg:  config.Config{TLSInsecureSkipVerify: true},
+	}
+
+	if err := HandleFileUploadHTTP(context.Background(), env, "cmd-http-upload-mismatch", destPath, ts.URL+"/file", int64(len(data)+1)); err != nil {
+		t.Fatalf("HandleFileUploadHTTP failed: %v", err)
+	}
+
+	if len(writer.msgs) != 1 {
+		t.Fatalf("expected 1 message, got %d", len(writer.msgs))
+	}
+	var result wire.CommandResult
+	if err := msgpack.Unmarshal(writer.msgs[0], &result); err != nil {
+		t.Fatalf("failed to unmarshal command result: %v", err)
+	}
+	if result.OK {
+		t.Fatalf("expected OK=false for size mismatch")
+	}
+}