feat: add agent-friendly ufile CLI
This commit is contained in:
loveuer
@@ -0,0 +1,301 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"runtime"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func runTestCLI(t *testing.T, args []string, input string) (int, string, string) {
|
||||
t.Helper()
|
||||
var stdout, stderr bytes.Buffer
|
||||
code := run(context.Background(), args, streams{in: strings.NewReader(input), out: &stdout, errOut: &stderr})
|
||||
return code, stdout.String(), stderr.String()
|
||||
}
|
||||
|
||||
func TestHelpIsAvailableWithoutConfiguration(t *testing.T) {
|
||||
t.Setenv("UFILE_HOST", "")
|
||||
t.Setenv("UFILE_SERVER", "")
|
||||
code, stdout, stderr := runTestCLI(t, []string{"download", "--help"}, "")
|
||||
if code != exitOK || !strings.Contains(stdout, "-O, --output") || !strings.Contains(stdout, "达到 5 次") || stderr != "" {
|
||||
t.Fatalf("code=%d, stdout=%q, stderr=%q", code, stdout, stderr)
|
||||
}
|
||||
}
|
||||
|
||||
func TestInvalidFlagWritesOnlyJSONError(t *testing.T) {
|
||||
code, stdout, stderr := runTestCLI(t, []string{"--not-a-real-option"}, "")
|
||||
if code != exitUsage || stdout != "" || strings.Count(strings.TrimSpace(stderr), "\n") != 0 {
|
||||
t.Fatalf("code=%d, stdout=%q, stderr=%q", code, stdout, stderr)
|
||||
}
|
||||
var payload map[string]any
|
||||
if err := json.Unmarshal([]byte(stderr), &payload); err != nil {
|
||||
t.Fatalf("stderr is not JSON: %v; %q", err, stderr)
|
||||
}
|
||||
}
|
||||
|
||||
func TestUnsupportedShortAliasesAreRejected(t *testing.T) {
|
||||
tests := [][]string{
|
||||
{"-h"},
|
||||
{"-c", "config.json", "version"},
|
||||
{"-t", "1m", "version"},
|
||||
{"-H", "https://example.com", "auth", "login", "-P"},
|
||||
{"-H", "https://example.com", "auth", "login", "-u", "agent"},
|
||||
{"-H", "https://example.com", "upload", "file.txt", "-N"},
|
||||
{"-H", "https://example.com", "upload", "file.txt", "-r", "1h"},
|
||||
{"-H", "https://example.com", "upload", "file.txt", "-T", "text/plain"},
|
||||
{"-H", "https://example.com", "download", "Ab3dE8xP", "-c", "123456", "-O", "out"},
|
||||
{"-H", "https://example.com", "download", "Ab3dE8xP", "-O", "out", "-f"},
|
||||
}
|
||||
for _, args := range tests {
|
||||
code, stdout, stderr := runTestCLI(t, args, "")
|
||||
if code != exitUsage || stdout != "" || !strings.Contains(stderr, `"code":"usage_error"`) {
|
||||
t.Fatalf("args=%q, code=%d, stdout=%q, stderr=%q", args, code, stdout, stderr)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestAuthLoginStoresCookieWithoutPassword(t *testing.T) {
|
||||
expires := time.Now().Add(30 * 24 * time.Hour).UTC().Truncate(time.Second)
|
||||
server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
|
||||
if request.URL.Path != "/api/auth/login" {
|
||||
t.Errorf("path = %s", request.URL.Path)
|
||||
}
|
||||
var body map[string]string
|
||||
if err := json.NewDecoder(request.Body).Decode(&body); err != nil {
|
||||
t.Error(err)
|
||||
}
|
||||
if body["username"] != "agent" || body["password"] != "secret-value" {
|
||||
t.Errorf("body = %#v", body)
|
||||
}
|
||||
http.SetCookie(writer, &http.Cookie{Name: sessionCookieName, Value: "session-secret", Expires: expires, Path: "/"})
|
||||
writer.Header().Set("Content-Type", "application/json")
|
||||
_, _ = io.WriteString(writer, `{"id":2,"username":"agent","canUpload":true}`)
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
configPath := filepath.Join(t.TempDir(), "nested", "config.json")
|
||||
code, stdout, stderr := runTestCLI(t, []string{"-H", server.URL, "--config", configPath, "auth", "login", "-U", "agent", "--password-stdin"}, "secret-value\n")
|
||||
if code != exitOK || stderr != "" {
|
||||
t.Fatalf("code=%d, stdout=%s, stderr=%s", code, stdout, stderr)
|
||||
}
|
||||
config, err := loadConfig(configPath)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if config.Host != server.URL || config.Session != "session-secret" || config.ExpiresAt.IsZero() {
|
||||
t.Fatalf("config = %#v", config)
|
||||
}
|
||||
data, err := os.ReadFile(configPath)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if strings.Contains(string(data), "secret-value") {
|
||||
t.Fatal("password was written to config")
|
||||
}
|
||||
if runtime.GOOS != "windows" {
|
||||
info, _ := os.Stat(configPath)
|
||||
if info.Mode().Perm() != 0o600 {
|
||||
t.Fatalf("config mode = %o", info.Mode().Perm())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestAuthWhoamiAndLogoutUseStoredSession(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
|
||||
cookie, err := request.Cookie(sessionCookieName)
|
||||
if err != nil || cookie.Value != "stored-session" {
|
||||
t.Errorf("cookie=%#v, err=%v", cookie, err)
|
||||
}
|
||||
switch {
|
||||
case request.Method == http.MethodGet && request.URL.Path == "/api/auth/me":
|
||||
writer.Header().Set("Content-Type", "application/json")
|
||||
_, _ = io.WriteString(writer, `{"id":2,"username":"agent","canUpload":true}`)
|
||||
case request.Method == http.MethodPost && request.URL.Path == "/api/auth/logout":
|
||||
writer.WriteHeader(http.StatusNoContent)
|
||||
default:
|
||||
http.NotFound(writer, request)
|
||||
}
|
||||
}))
|
||||
defer server.Close()
|
||||
configPath := filepath.Join(t.TempDir(), "config.json")
|
||||
if err := saveConfig(configPath, configFile{Host: server.URL, Session: "stored-session"}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
code, stdout, stderr := runTestCLI(t, []string{"--config", configPath, "auth", "whoami"}, "")
|
||||
if code != exitOK || stderr != "" || !strings.Contains(stdout, `"username":"agent"`) {
|
||||
t.Fatalf("whoami code=%d, stdout=%s, stderr=%s", code, stdout, stderr)
|
||||
}
|
||||
code, stdout, stderr = runTestCLI(t, []string{"--config", configPath, "auth", "logout"}, "")
|
||||
if code != exitOK || stderr != "" || !strings.Contains(stdout, `"loggedOut":true`) {
|
||||
t.Fatalf("logout code=%d, stdout=%s, stderr=%s", code, stdout, stderr)
|
||||
}
|
||||
config, err := loadConfig(configPath)
|
||||
if err != nil || config.Session != "" {
|
||||
t.Fatalf("config=%#v, err=%v", config, err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestUploadSupportsShortOptionsAfterPathAndGeneratesCode(t *testing.T) {
|
||||
t.Setenv("UFILE_SESSION", "test-session")
|
||||
var receivedCode string
|
||||
server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
|
||||
if cookie, err := request.Cookie(sessionCookieName); err != nil || cookie.Value != "test-session" {
|
||||
t.Errorf("cookie = %#v, err = %v", cookie, err)
|
||||
}
|
||||
if err := request.ParseMultipartForm(2 << 20); err != nil {
|
||||
t.Error(err)
|
||||
}
|
||||
if request.FormValue("retention") != "8h" || request.FormValue("downloads") != "3" {
|
||||
t.Errorf("form = %#v", request.Form)
|
||||
}
|
||||
receivedCode = request.FormValue("pickupCode")
|
||||
if !pickupCodePattern.MatchString(receivedCode) {
|
||||
t.Errorf("pickup code = %q", receivedCode)
|
||||
}
|
||||
file, header, err := request.FormFile("file")
|
||||
if err != nil {
|
||||
t.Error(err)
|
||||
} else {
|
||||
if header.Header.Get("Content-Type") != "text/plain; charset=utf-8" {
|
||||
t.Errorf("content type = %q", header.Header.Get("Content-Type"))
|
||||
}
|
||||
content, _ := io.ReadAll(file)
|
||||
file.Close()
|
||||
if string(content) != "agent artifact" {
|
||||
t.Errorf("content = %q", content)
|
||||
}
|
||||
}
|
||||
writer.Header().Set("Content-Type", "application/json")
|
||||
writer.WriteHeader(http.StatusCreated)
|
||||
_, _ = io.WriteString(writer, `{"id":1,"hash":"Ab3dE8xP","name":"artifact.txt","size":14,"expiresAt":null,"maxDownloads":3}`)
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
path := filepath.Join(t.TempDir(), "artifact.txt")
|
||||
if err := os.WriteFile(path, []byte("agent artifact"), 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
code, stdout, stderr := runTestCLI(t, []string{"-H", server.URL, "upload", path, "-R", "8h", "-D", "3"}, "")
|
||||
if code != exitOK || stderr != "" {
|
||||
t.Fatalf("code=%d, stdout=%s, stderr=%s", code, stdout, stderr)
|
||||
}
|
||||
var result uploadResult
|
||||
if err := json.Unmarshal([]byte(stdout), &result); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if result.PickupCode != receivedCode || !strings.Contains(result.DownloadURL, "code="+receivedCode) || !strings.Contains(result.PreviewURL, "preview=true") {
|
||||
t.Fatalf("result = %#v", result)
|
||||
}
|
||||
}
|
||||
|
||||
func TestUploadFromStdinCanBeExplicitlyPublic(t *testing.T) {
|
||||
t.Setenv("UFILE_SESSION", "test-session")
|
||||
server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
|
||||
if err := request.ParseMultipartForm(2 << 20); err != nil {
|
||||
t.Error(err)
|
||||
}
|
||||
if code := request.FormValue("pickupCode"); code != "" {
|
||||
t.Errorf("pickup code = %q", code)
|
||||
}
|
||||
file, header, err := request.FormFile("file")
|
||||
if err != nil {
|
||||
t.Error(err)
|
||||
} else {
|
||||
content, _ := io.ReadAll(file)
|
||||
file.Close()
|
||||
if header.Filename != "agent.txt" || string(content) != "stdin artifact" {
|
||||
t.Errorf("filename=%q, content=%q", header.Filename, content)
|
||||
}
|
||||
}
|
||||
writer.Header().Set("Content-Type", "application/json")
|
||||
writer.WriteHeader(http.StatusCreated)
|
||||
_, _ = io.WriteString(writer, `{"hash":"Cd4fG6hJ","name":"agent.txt","size":14,"maxDownloads":5}`)
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
code, stdout, stderr := runTestCLI(t, []string{"-H", server.URL, "upload", "-", "-N", "agent.txt", "--no-pickup-code"}, "stdin artifact")
|
||||
if code != exitOK || stderr != "" || strings.Contains(stdout, "pickupCode") || !strings.Contains(stdout, `"downloadUrl"`) {
|
||||
t.Fatalf("code=%d, stdout=%s, stderr=%s", code, stdout, stderr)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDownloadUsesOneRequestAndWritesAtomically(t *testing.T) {
|
||||
var requests atomic.Int32
|
||||
server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
|
||||
requests.Add(1)
|
||||
if request.URL.Path != "/uf/Ab3dE8xP" || request.URL.Query().Get("code") != "123456" {
|
||||
t.Errorf("url = %s", request.URL.String())
|
||||
}
|
||||
writer.Header().Set("Content-Type", "text/plain")
|
||||
writer.Header().Set("Content-Disposition", `attachment; filename="result.txt"`)
|
||||
_, _ = io.WriteString(writer, "downloaded")
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
output := filepath.Join(t.TempDir(), "result.txt")
|
||||
code, stdout, stderr := runTestCLI(t, []string{"--server", server.URL, "download", "Ab3dE8xP", "-C", "123456", "-O", output}, "")
|
||||
if code != exitOK || stderr != "" || requests.Load() != 1 {
|
||||
t.Fatalf("code=%d, requests=%d, stdout=%s, stderr=%s", code, requests.Load(), stdout, stderr)
|
||||
}
|
||||
content, err := os.ReadFile(output)
|
||||
if err != nil || string(content) != "downloaded" {
|
||||
t.Fatalf("content=%q, err=%v", content, err)
|
||||
}
|
||||
var result downloadResult
|
||||
if err := json.Unmarshal([]byte(stdout), &result); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if result.Name != "result.txt" || result.Size != int64(len(content)) {
|
||||
t.Fatalf("result = %#v", result)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDownloadDoesNotConsumeAccessWhenTargetExists(t *testing.T) {
|
||||
var requests atomic.Int32
|
||||
server := httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) { requests.Add(1) }))
|
||||
defer server.Close()
|
||||
output := filepath.Join(t.TempDir(), "exists.txt")
|
||||
if err := os.WriteFile(output, []byte("keep"), 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
code, _, stderr := runTestCLI(t, []string{"-H", server.URL, "download", "Ab3dE8xP", "-O", output}, "")
|
||||
if code != exitLocalIO || requests.Load() != 0 || !strings.Contains(stderr, "local_io_error") {
|
||||
t.Fatalf("code=%d, requests=%d, stderr=%s", code, requests.Load(), stderr)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDownloadToStdoutContainsOnlyFileBytes(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
|
||||
_, _ = io.WriteString(writer, "raw bytes")
|
||||
}))
|
||||
defer server.Close()
|
||||
code, stdout, stderr := runTestCLI(t, []string{"-H", server.URL, "download", "Ab3dE8xP", "-O", "-"}, "")
|
||||
if code != exitOK || stdout != "raw bytes" || stderr != "" {
|
||||
t.Fatalf("code=%d, stdout=%q, stderr=%q", code, stdout, stderr)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRemoteErrorsAreMachineReadable(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
|
||||
writer.Header().Set("Content-Type", "application/json")
|
||||
writer.WriteHeader(http.StatusNotFound)
|
||||
_, _ = io.WriteString(writer, `{"error":"文件不存在或取件码错误","code":"file_unavailable"}`)
|
||||
}))
|
||||
defer server.Close()
|
||||
output := filepath.Join(t.TempDir(), "missing.txt")
|
||||
code, stdout, stderr := runTestCLI(t, []string{"-H", server.URL, "download", "Ab3dE8xP", "-O", output}, "")
|
||||
if code != exitUnavailable || stdout != "" || !strings.Contains(stderr, `"code":"file_unavailable"`) || !strings.Contains(stderr, `"httpStatus":404`) {
|
||||
t.Fatalf("code=%d, stdout=%s, stderr=%s", code, stdout, stderr)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user