twins-upstream/serve_file.go

package main

import (
	"bytes"
	"fmt"
	"io"
	"net"
	"net/url"
	"os"
	"path/filepath"
	"sort"

	"github.com/h2non/filetype"
)

func buildDirList(request *url.URL, dirPath string) ([]byte, error) {
	var (
		b        = &bytes.Buffer{}
		files    []os.FileInfo
		numDirs  int
		numFiles int
	)
	err := filepath.Walk(dirPath, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		} else if path == dirPath {
			return nil
		}
		files = append(files, info)
		if info.IsDir() || info.Mode()&os.ModeSymlink != 0 {
			numDirs++
		} else {
			numFiles++
		}
		if info.IsDir() {
			return filepath.SkipDir
		}
		return nil
	})
	if err != nil {
		return nil, err
	}
	// List directories first
	sort.Slice(files, func(i, j int) bool {
		iDir := files[i].IsDir() || files[i].Mode()&os.ModeSymlink != 0
		jDir := files[j].IsDir() || files[j].Mode()&os.ModeSymlink != 0
		if iDir != jDir {
			return iDir
		}
		return i < j
	})

	fmt.Fprintf(b, "# %s%s", request.Path, newLine)
	if numDirs == 1 {
		b.Write([]byte("1 directory"))
	} else {
		fmt.Fprintf(b, "%d directories", numDirs)
	}
	b.Write([]byte(", "))
	if numDirs == 1 {
		b.Write([]byte("1 file"))
	} else {
		fmt.Fprintf(b, "%d files", numFiles)
	}
	b.Write([]byte(newLine + newLine))

	if request.Path != "/" {
		b.Write([]byte("=> ../ ../" + newLine + newLine))
	}

	for _, info := range files {
		fileName := info.Name()
		filePath := url.PathEscape(info.Name())
		if info.IsDir() || info.Mode()&os.ModeSymlink != 0 {
			fileName += "/"
			filePath += "/"
		}

		b.Write([]byte("=> " + fileName + " " + filePath + newLine))

		if info.IsDir() || info.Mode()&os.ModeSymlink != 0 {
			b.Write([]byte(newLine))
			continue
		}

		modified := "Never"
		if !info.ModTime().IsZero() {
			modified = info.ModTime().Format("2006-01-02 3:04 PM")
		}
		b.Write([]byte(modified + " - " + formatFileSize(info.Size()) + newLine + newLine))
	}

	return b.Bytes(), nil
}

func serveDirList(c net.Conn, serve *pathConfig, request *url.URL, dirPath string) int {
	dirList, err := buildDirList(request, dirPath)
	if err != nil {
		return writeStatus(c, statusTemporaryFailure)
	}

	writeSuccess(c, serve, geminiType, -1)
	c.Write(dirList)
	return statusSuccess
}

func serveFile(c net.Conn, serve *pathConfig, filePath string) {
	// Open file
	file, _ := os.Open(filePath)
	defer file.Close()

	// Read content type
	var (
		buf = make([]byte, 261)
		n   int
	)
	contentType := config.Types[filepath.Ext(filePath)]
	if contentType == "" {
		n, _ = file.Read(buf)
		kind, err := filetype.Match(buf[:n])
		if err == nil && kind != filetype.Unknown && kind.MIME.Value != "" {
			contentType = kind.MIME.Value
		} else {
			contentType = plainType
		}
	}

	// Read file size
	size := int64(-1)
	info, err := file.Stat()
	if err == nil {
		size = info.Size()
	}

	// Write response header
	writeSuccess(c, serve, contentType, size)

	// Write file contents
	if n > 0 {
		c.Write(buf[:n])
	}
	io.Copy(c, file)
}