1package pastes
  2
  3import (
  4	"context"
  5	"errors"
  6	"fmt"
  7	"io"
  8	"log/slog"
  9	"strings"
 10	"testing"
 11	"testing/iotest"
 12	"time"
 13
 14	"github.com/picosh/pico/pkg/db"
 15	"github.com/picosh/pico/pkg/db/stub"
 16	"github.com/picosh/pico/pkg/pssh"
 17	sendutils "github.com/picosh/pico/pkg/send/utils"
 18	"golang.org/x/crypto/ssh"
 19)
 20
 21type mockFeatureDB struct {
 22	*stub.StubDB
 23	features map[string]*db.FeatureFlag
 24}
 25
 26func newMockFeatureDB() *mockFeatureDB {
 27	return &mockFeatureDB{
 28		StubDB:   stub.NewStubDB(slog.Default()),
 29		features: make(map[string]*db.FeatureFlag),
 30	}
 31}
 32
 33func (m *mockFeatureDB) FindFeature(userID, name string) (*db.FeatureFlag, error) {
 34	if ff, ok := m.features[userID+":"+name]; ok {
 35		return ff, nil
 36	}
 37	return nil, fmt.Errorf("feature flag %s not found for user %s", name, userID)
 38}
 39
 40func (m *mockFeatureDB) setFeature(userID, name string, expiresAt time.Time) {
 41	m.features[userID+":"+name] = &db.FeatureFlag{UserID: userID, Name: name, ExpiresAt: &expiresAt}
 42}
 43
 44func TestUploadHandlerValidate(t *testing.T) {
 45	cfg := NewConfig("pastes-test")
 46	userID := "user-1"
 47	future := time.Now().Add(24 * time.Hour)
 48	past := time.Now().Add(-24 * time.Hour)
 49
 50	createSession := func(user *db.User) *pssh.SSHServerConnSession {
 51		conn := &pssh.SSHServerConn{
 52			Conn: &ssh.ServerConn{
 53				Permissions: &ssh.Permissions{
 54					Extensions: map[string]string{},
 55				},
 56			},
 57			Logger: slog.Default(),
 58		}
 59		sesh := &pssh.SSHServerConnSession{
 60			SSHServerConn: conn,
 61			Ctx:           context.Background(),
 62		}
 63		if user != nil {
 64			pssh.SetUser(sesh, user)
 65		}
 66		return sesh
 67	}
 68
 69	tests := []struct {
 70		name    string
 71		setup   func(m *mockFeatureDB)
 72		user    *db.User
 73		wantErr string
 74	}{
 75		{
 76			name:    "nil user",
 77			setup:   func(m *mockFeatureDB) {},
 78			wantErr: "could not get user",
 79		},
 80		{
 81			name:    "no invite or pico+",
 82			setup:   func(m *mockFeatureDB) {},
 83			user:    &db.User{ID: userID, Name: "tester"},
 84			wantErr: "requires an invitation or pico+",
 85		},
 86		{
 87			name:  "pastes feature",
 88			setup: func(m *mockFeatureDB) { m.setFeature(userID, "pastes", future) },
 89			user:  &db.User{ID: userID, Name: "tester"},
 90		},
 91		{
 92			name:  "pico+ feature",
 93			setup: func(m *mockFeatureDB) { m.setFeature(userID, "plus", future) },
 94			user:  &db.User{ID: userID, Name: "tester"},
 95		},
 96		{
 97			name:  "expired pico+ with active pastes",
 98			setup: func(m *mockFeatureDB) { m.setFeature(userID, "plus", past); m.setFeature(userID, "pastes", future) },
 99			user:  &db.User{ID: userID, Name: "tester"},
100		},
101		{
102			name:    "expired pastes",
103			setup:   func(m *mockFeatureDB) { m.setFeature(userID, "pastes", past) },
104			user:    &db.User{ID: userID, Name: "tester"},
105			wantErr: "pastes access has expired",
106		},
107		{
108			name:    "expired pico+",
109			setup:   func(m *mockFeatureDB) { m.setFeature(userID, "plus", past) },
110			user:    &db.User{ID: userID, Name: "tester"},
111			wantErr: "pico+ has expired",
112		},
113	}
114
115	for _, tt := range tests {
116		t.Run(tt.name, func(t *testing.T) {
117			mockDB := newMockFeatureDB()
118			tt.setup(mockDB)
119			handler := NewUploadHandler(cfg, mockDB)
120
121			err := handler.Validate(createSession(tt.user))
122			if tt.wantErr == "" {
123				if err != nil {
124					t.Fatalf("unexpected error: %v", err)
125				}
126				return
127			}
128			if err == nil || !strings.Contains(err.Error(), tt.wantErr) {
129				t.Fatalf("expected error containing %q, got %v", tt.wantErr, err)
130			}
131		})
132	}
133}
134
135func TestSetMetaKeepsSettings(t *testing.T) {
136	soon := time.Now().Add(time.Hour)
137	past := time.Now().Add(-time.Hour)
138	// Pastes expire after calendar days, which a daylight saving change
139	// makes an hour longer or shorter.
140	defaultExpiry := time.Until(time.Now().AddDate(0, 0, DEFAULT_EXPIRES_AT)).Round(time.Minute)
141	tests := []struct {
142		name   string
143		cmd    []string
144		cur    *db.Post
145		hidden bool
146		// expires is how long until the paste expires, or 0 for never.
147		expires time.Duration
148	}{
149		{name: "new", expires: defaultExpiry},
150		{name: "new with arguments", cmd: []string{"hidden=true", "expires=2h"}, hidden: true, expires: 2 * time.Hour},
151		{name: "existing", cur: &db.Post{Hidden: true, ExpiresAt: &soon}, hidden: true, expires: time.Hour},
152		{name: "existing without expiry", cur: &db.Post{Hidden: true}, hidden: true},
153		{name: "existing past its expiry", cur: &db.Post{Hidden: true, ExpiresAt: &past}, hidden: true, expires: defaultExpiry},
154		{name: "existing with arguments", cmd: []string{"hidden=false", "expires=true"}, cur: &db.Post{Hidden: true}, expires: defaultExpiry},
155		{name: "existing made permanent", cmd: []string{"expires=false"}, cur: &db.Post{ExpiresAt: &soon}},
156	}
157	for _, tt := range tests {
158		t.Run(tt.name, func(t *testing.T) {
159			s := &pssh.SSHServerConnSession{Ctx: context.Background()}
160			s.SetValue("command", tt.cmd)
161			data := &db.Post{Filename: "a.txt", Slug: "a"}
162			setMeta(s, tt.cur, data)
163			if data.Hidden != tt.hidden {
164				t.Errorf("hidden = %v, want %v", data.Hidden, tt.hidden)
165			}
166			switch {
167			case tt.expires == 0 && data.ExpiresAt != nil:
168				t.Errorf("expires at %v, want never", data.ExpiresAt)
169			case tt.expires != 0 && (data.ExpiresAt == nil || time.Until(*data.ExpiresAt).Round(time.Minute) != tt.expires):
170				t.Errorf("expires at %v, want in %v", data.ExpiresAt, tt.expires)
171			}
172		})
173	}
174}
175
176func TestWriteReadError(t *testing.T) {
177	checksumErr := errors.New("rsync: file failed verification")
178	tests := []struct {
179		name string
180		text string
181		err  error
182	}{
183		{name: "before contents", err: io.ErrClosedPipe},
184		{name: "partial contents", text: "partial upload", err: io.ErrUnexpectedEOF},
185		{name: "failed verification", text: "complete upload", err: checksumErr},
186	}
187	for _, tt := range tests {
188		t.Run(tt.name, func(t *testing.T) {
189			logger := slog.New(slog.DiscardHandler)
190			session := &pssh.SSHServerConnSession{
191				SSHServerConn: &pssh.SSHServerConn{
192					Conn:   &ssh.ServerConn{Permissions: &ssh.Permissions{Extensions: map[string]string{}}},
193					Logger: logger,
194				},
195				Ctx: context.Background(),
196			}
197			user := &db.User{ID: "user-1", Name: "tester"}
198			pssh.SetUser(session, user)
199			posts := stub.NewPostsDB(logger, user)
200			handler := NewUploadHandler(&Config{Space: "pastes", MaxAssetSize: 1024}, posts)
201			reader := io.MultiReader(strings.NewReader(tt.text), iotest.ErrReader(tt.err))
202
203			message, err := handler.Write(session, &sendutils.FileEntry{Filepath: "/test.txt", Reader: reader})
204			if !errors.Is(err, tt.err) {
205				t.Fatalf("error = %v, want %v", err, tt.err)
206			}
207			if message != "" {
208				t.Fatalf("failed upload returned a success message: %q", message)
209			}
210			if all, _ := posts.FindAllPostsByUser(user.ID, "pastes"); len(all) != 0 {
211				t.Fatalf("failed upload stored %d posts", len(all))
212			}
213		})
214	}
215}