- StdinSource reads from an injectable io.Reader (INPUT-03) - URLSource fetches http/https with 30s timeout, 50MB cap, scheme whitelist, and Content-Type filter (INPUT-04) - ClipboardSource wraps atotto/clipboard with graceful fallback for missing tooling (INPUT-05) - emitByteChunks local helper mirrors file.go windowing to stay independent of sibling wave-1 plans - Tests cover happy path, cancellation, redirects, oversize bodies, binary content types, scheme rejection, and clipboard error paths
103 lines
3.0 KiB
Go
103 lines
3.0 KiB
Go
package sources
|
|
|
|
import (
|
|
"context"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"strings"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/stretchr/testify/require"
|
|
|
|
"github.com/salvacybersec/keyhunter/pkg/types"
|
|
)
|
|
|
|
func drainURL(t *testing.T, src Source) ([]types.Chunk, error) {
|
|
t.Helper()
|
|
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
|
|
defer cancel()
|
|
out := make(chan types.Chunk, 65536)
|
|
errCh := make(chan error, 1)
|
|
go func() { errCh <- src.Chunks(ctx, out); close(out) }()
|
|
var got []types.Chunk
|
|
for c := range out {
|
|
got = append(got, c)
|
|
}
|
|
return got, <-errCh
|
|
}
|
|
|
|
func TestURLSource_Fetches(t *testing.T) {
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
w.Header().Set("Content-Type", "text/plain")
|
|
_, _ = w.Write([]byte("API_KEY=sk-live-xyz"))
|
|
}))
|
|
defer srv.Close()
|
|
|
|
chunks, err := drainURL(t, NewURLSource(srv.URL))
|
|
require.NoError(t, err)
|
|
require.Len(t, chunks, 1)
|
|
require.Equal(t, "url:"+srv.URL, chunks[0].Source)
|
|
require.Equal(t, "API_KEY=sk-live-xyz", string(chunks[0].Data))
|
|
}
|
|
|
|
func TestURLSource_RejectsBinaryContentType(t *testing.T) {
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
w.Header().Set("Content-Type", "image/png")
|
|
_, _ = w.Write([]byte{0x89, 0x50, 0x4e, 0x47})
|
|
}))
|
|
defer srv.Close()
|
|
|
|
_, err := drainURL(t, NewURLSource(srv.URL))
|
|
require.Error(t, err)
|
|
require.Contains(t, err.Error(), "Content-Type")
|
|
}
|
|
|
|
func TestURLSource_RejectsNonHTTPScheme(t *testing.T) {
|
|
_, err := drainURL(t, NewURLSource("file:///etc/passwd"))
|
|
require.Error(t, err)
|
|
require.Contains(t, err.Error(), "unsupported scheme")
|
|
}
|
|
|
|
func TestURLSource_Rejects500(t *testing.T) {
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
http.Error(w, "boom", http.StatusInternalServerError)
|
|
}))
|
|
defer srv.Close()
|
|
|
|
_, err := drainURL(t, NewURLSource(srv.URL))
|
|
require.Error(t, err)
|
|
require.Contains(t, err.Error(), "500")
|
|
}
|
|
|
|
func TestURLSource_RejectsOversizeBody(t *testing.T) {
|
|
// Serve body just over the cap.
|
|
big := strings.Repeat("a", int(MaxURLContentLength)+10)
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
w.Header().Set("Content-Type", "text/plain")
|
|
_, _ = w.Write([]byte(big))
|
|
}))
|
|
defer srv.Close()
|
|
|
|
_, err := drainURL(t, NewURLSource(srv.URL))
|
|
require.Error(t, err)
|
|
}
|
|
|
|
func TestURLSource_FollowsRedirect(t *testing.T) {
|
|
target := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
w.Header().Set("Content-Type", "text/plain")
|
|
_, _ = w.Write([]byte("redirected body"))
|
|
}))
|
|
defer target.Close()
|
|
|
|
redirector := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
http.Redirect(w, r, target.URL, http.StatusMovedPermanently)
|
|
}))
|
|
defer redirector.Close()
|
|
|
|
chunks, err := drainURL(t, NewURLSource(redirector.URL))
|
|
require.NoError(t, err)
|
|
require.NotEmpty(t, chunks)
|
|
require.Contains(t, string(chunks[0].Data), "redirected body")
|
|
}
|