go-omni/examples/signedRead/main.go

// signedRead is an offline, loopback-only preview/download demonstration.
package main

import (
	"bytes"
	"context"
	"crypto/rand"
	"errors"
	"flag"
	"fmt"
	"html/template"
	"io"
	"mime"
	"net"
	"net/http"
	"os"
	"os/signal"
	"strconv"
	"time"

	"git.yuntuo.hk/k2safe/go-omni/storage"
	"git.yuntuo.hk/k2safe/go-omni/storage/local"
)

func main() {
	check := flag.Bool("check", false, "run HTTP checks and exit")
	flag.Parse()
	if err := run(*check); err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}
}

func run(check bool) error {
	root, err := os.MkdirTemp("", "go-omni-signed-read-")
	if err != nil {
		return err
	}
	defer os.RemoveAll(root)
	files, err := local.New(root, 1<<20)
	if err != nil {
		return err
	}
	content := samplePDF()
	if _, err = files.Put(context.Background(), "sample.pdf", bytes.NewReader(content), int64(len(content)), "application/pdf"); err != nil {
		return err
	}
	listener, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		return err
	}
	defer listener.Close()
	baseURL := "http://" + listener.Addr().String()
	secret := make([]byte, 32)
	if _, err = rand.Read(secret); err != nil {
		return err
	}
	signer, err := storage.NewSigner(files, baseURL+"/file", secret)
	if err != nil {
		return err
	}
	// Production applications authorize the caller before this operation. The
	// sample intentionally exposes only a generated fixture on loopback.
	preview, err := signer.PresignGetWithOptions(context.Background(), "sample.pdf", 30*time.Minute, storage.ReadOptions{
		Disposition: storage.DispositionInline, FileName: "示例报告.pdf", ContentType: "application/pdf",
	})
	if err != nil {
		return err
	}
	download, err := signer.PresignGetWithOptions(context.Background(), "sample.pdf", 30*time.Minute, storage.ReadOptions{FileName: "示例报告.pdf"})
	if err != nil {
		return err
	}
	page := template.Must(template.New("page").Parse(`<!doctype html><html lang="zh-CN"><meta charset="utf-8"><title>Signed read example</title>
<body><h1>文件预览与下载</h1><p><a href="{{.Preview.URL}}">打开示例 PDF</a> · <a href="{{.Download.URL}}">下载示例 PDF</a></p>
<p>链接到期时间：{{.Preview.ExpiresAt}}</p><p>只读演示：无数据库、业务身份或公开磁盘目录。</p></body></html>`))
	mux := http.NewServeMux()
	mux.Handle("/file", signer)
	mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path != "/" || r.Method != http.MethodGet {
			http.NotFound(w, r)
			return
		}
		w.Header().Set("Content-Type", "text/html; charset=utf-8")
		w.Header().Set("Cache-Control", "private, no-store")
		w.Header().Set("Referrer-Policy", "no-referrer")
		_ = page.Execute(w, struct{ Preview, Download storage.ReadLink }{preview, download})
	})
	server := &http.Server{Handler: mux, ReadHeaderTimeout: 5 * time.Second, IdleTimeout: 30 * time.Second}
	done := make(chan error, 1)
	go func() { done <- server.Serve(listener) }()
	defer server.Close()
	if check {
		if err := checkHTTP(preview, download, content); err != nil {
			return err
		}
		fmt.Println("signed read: PDF preview, Chinese filename, download, HEAD and Range passed")
		return nil
	}
	// Never print signed URLs or enable request logging containing query strings.
	fmt.Println("Open " + baseURL + " (Ctrl+C to stop)")
	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt)
	defer stop()
	select {
	case err := <-done:
		if errors.Is(err, http.ErrServerClosed) {
			return nil
		}
		return err
	case <-ctx.Done():
		shutdown, cancel := context.WithTimeout(context.Background(), 5*time.Second)
		defer cancel()
		return server.Shutdown(shutdown)
	}
}

func checkHTTP(preview, download storage.ReadLink, content []byte) error {
	client := &http.Client{Timeout: 5 * time.Second}
	for _, tc := range []struct {
		link                      storage.ReadLink
		method, span, disposition string
		status                    int
		body                      []byte
	}{
		{preview, "GET", "", "inline", 200, content},
		{preview, "HEAD", "", "inline", 200, nil},
		{preview, "GET", "bytes=0-7", "inline", 206, content[:8]},
		{download, "GET", "", "attachment", 200, content},
	} {
		req, err := http.NewRequest(tc.method, tc.link.URL, nil)
		if err != nil {
			return errors.New("invalid example link")
		}
		if tc.span != "" {
			req.Header.Set("Range", tc.span)
		}
		response, err := client.Do(req)
		if err != nil {
			return errors.New("example request failed")
		}
		body, err := io.ReadAll(response.Body)
		response.Body.Close()
		kind, params, headerErr := mime.ParseMediaType(response.Header.Get("Content-Disposition"))
		if err != nil || headerErr != nil || response.StatusCode != tc.status || !bytes.Equal(body, tc.body) || kind != tc.disposition || params["filename"] != "示例报告.pdf" || response.Header.Get("Content-Type") != "application/pdf" {
			return errors.New("example response did not match signed options")
		}
		size := len(tc.body)
		if tc.method == "HEAD" {
			size = len(content)
		}
		if response.Header.Get("Content-Length") != strconv.Itoa(size) || response.Header.Get("Cache-Control") != "private, no-store" {
			return errors.New("example size or cache policy mismatch")
		}
	}
	return nil
}

// A complete, single-page PDF with a real cross-reference table. No generator
// dependency or business document is needed for this example.
func samplePDF() []byte {
	stream := "BT /F1 24 Tf 72 720 Td (go-omni signed read) Tj ET\n"
	objects := []string{
		"<< /Type /Catalog /Pages 2 0 R >>",
		"<< /Type /Pages /Kids [3 0 R] /Count 1 >>",
		"<< /Type /Page /Parent 2 0 R /MediaBox [0 0 612 792] /Resources << /Font << /F1 4 0 R >> >> /Contents 5 0 R >>",
		"<< /Type /Font /Subtype /Type1 /BaseFont /Helvetica >>",
		fmt.Sprintf("<< /Length %d >>\nstream\n%sendstream", len(stream), stream),
	}
	var pdf bytes.Buffer
	pdf.WriteString("%PDF-1.4\n")
	offsets := []int{0}
	for i, object := range objects {
		offsets = append(offsets, pdf.Len())
		fmt.Fprintf(&pdf, "%d 0 obj\n%s\nendobj\n", i+1, object)
	}
	xref := pdf.Len()
	fmt.Fprintf(&pdf, "xref\n0 %d\n0000000000 65535 f \n", len(offsets))
	for _, offset := range offsets[1:] {
		fmt.Fprintf(&pdf, "%010d 00000 n \n", offset)
	}
	fmt.Fprintf(&pdf, "trailer\n<< /Size %d /Root 1 0 R >>\nstartxref\n%d\n%%%%EOF\n", len(offsets), xref)
	return pdf.Bytes()
}