58 lines
2.0 KiB
Go
58 lines
2.0 KiB
Go
package util
|
|
|
|
import (
|
|
"crypto/rand"
|
|
"github.com/stretchr/testify/require"
|
|
"net/http/httptest"
|
|
"testing"
|
|
)
|
|
|
|
func TestSniffWriter_WriteHTML(t *testing.T) {
|
|
rr := httptest.NewRecorder()
|
|
sw := NewContentTypeWriter(rr, "")
|
|
sw.Write([]byte("<script>alert('hi')</script>"))
|
|
require.Equal(t, "text/plain; charset=utf-8", rr.Header().Get("Content-Type"))
|
|
}
|
|
|
|
func TestSniffWriter_WriteTwoWriteCalls(t *testing.T) {
|
|
rr := httptest.NewRecorder()
|
|
sw := NewContentTypeWriter(rr, "")
|
|
sw.Write([]byte{0x25, 0x50, 0x44, 0x46, 0x2d, 0x11, 0x22, 0x33})
|
|
sw.Write([]byte("<script>alert('hi')</script>"))
|
|
require.Equal(t, "application/pdf", rr.Header().Get("Content-Type"))
|
|
}
|
|
|
|
func TestSniffWriter_NoSniffWriterWriteHTML(t *testing.T) {
|
|
// This test just makes sure that without the sniff-w, we would get text/html
|
|
|
|
rr := httptest.NewRecorder()
|
|
rr.Write([]byte("<script>alert('hi')</script>"))
|
|
require.Equal(t, "text/html; charset=utf-8", rr.Header().Get("Content-Type"))
|
|
}
|
|
|
|
func TestSniffWriter_WriteHTMLSplitIntoTwoWrites(t *testing.T) {
|
|
// This test shows how splitting the HTML into two Write() calls will still yield text/plain
|
|
|
|
rr := httptest.NewRecorder()
|
|
sw := NewContentTypeWriter(rr, "")
|
|
sw.Write([]byte("<scr"))
|
|
sw.Write([]byte("ipt>alert('hi')</script>"))
|
|
require.Equal(t, "text/plain; charset=utf-8", rr.Header().Get("Content-Type"))
|
|
}
|
|
|
|
func TestSniffWriter_WriteUnknownMimeType(t *testing.T) {
|
|
rr := httptest.NewRecorder()
|
|
sw := NewContentTypeWriter(rr, "")
|
|
randomBytes := make([]byte, 199)
|
|
rand.Read(randomBytes[5:]) // Start at an offset; the test kept failing randomly because it hit random magic strings
|
|
sw.Write(randomBytes)
|
|
require.Equal(t, "application/octet-stream", rr.Header().Get("Content-Type"))
|
|
}
|
|
|
|
func TestSniffWriter_WriteWithFilenameAPK(t *testing.T) {
|
|
rr := httptest.NewRecorder()
|
|
sw := NewContentTypeWriter(rr, "https://example.com/ntfy.apk")
|
|
sw.Write([]byte{0x50, 0x4B, 0x03, 0x04})
|
|
require.Equal(t, "application/vnd.android.package-archive", rr.Header().Get("Content-Type"))
|
|
}
|