pcap-web-tools/internal/web/split.go

package web

import (
	"archive/zip"
	"compress/flate"
	"errors"
	"fmt"
	"io"
	"log"
	"math"
	"net/http"
	"path"
	"strconv"
	"strings"
	"time"

	"pcap-web-tools/internal/pcap"
)

// maxSplitParts caps the number of files one split may produce.
var maxSplitParts = 10000

var sizeUnits = []struct {
	Name  string
	Bytes float64
}{
	{"KB", 1e3}, {"MB", 1e6}, {"GB", 1e9},
}

type splitForm struct {
	Mode    string // "packets" or "size"
	Packets string
	Size    string
	Unit    string
}

type splitPage struct {
	Form      splitForm
	Units     []string
	MaxUpload string
	MaxParts  int
	Error     string
}

func (s *Server) splitPage(f splitForm) splitPage {
	units := make([]string, len(sizeUnits))
	for i, u := range sizeUnits {
		units[i] = u.Name
	}
	return splitPage{
		Form:      f,
		Units:     units,
		MaxUpload: fmt.Sprintf("%d MB", s.cfg.MaxUploadBytes/1e6),
		MaxParts:  maxSplitParts,
	}
}

func (s *Server) handleSplitForm(w http.ResponseWriter, r *http.Request) {
	page := s.splitPage(splitForm{Mode: "size", Packets: "10000", Size: "10", Unit: "MB"})
	s.render(w, http.StatusOK, "split.html", "/split", page)
}

func (s *Server) handleSplit(w http.ResponseWriter, r *http.Request) {
	page := s.splitPage(splitForm{Mode: "size", Packets: "10000", Size: "10", Unit: "MB"})
	fail := func(status int, msg string) {
		page.Error = msg
		s.render(w, status, "split.html", "/split", page)
	}

	// Allow some slack over the file size limit for the multipart framing and form fields.
	r.Body = http.MaxBytesReader(w, r.Body, s.cfg.MaxUploadBytes+64<<10)
	// Uploads bigger than 1 MB are spooled to temporary files rather than kept in memory.
	if err := r.ParseMultipartForm(1 << 20); err != nil {
		var tooBig *http.MaxBytesError
		if errors.As(err, &tooBig) {
			fail(http.StatusRequestEntityTooLarge, "file too large; the maximum upload is "+page.MaxUpload)
		} else {
			fail(http.StatusBadRequest, "bad upload: "+err.Error())
		}
		return
	}
	defer r.MultipartForm.RemoveAll()

	page.Form = splitForm{
		Mode:    r.FormValue("mode"),
		Packets: r.FormValue("packets"),
		Size:    r.FormValue("size"),
		Unit:    r.FormValue("unit"),
	}
	limit, err := parseSplitLimit(page.Form)
	if err != nil {
		fail(http.StatusBadRequest, err.Error())
		return
	}
	file, hdr, err := r.FormFile("file")
	if err != nil {
		fail(http.StatusBadRequest, "choose a pcap file to split")
		return
	}
	defer file.Close()

	// First pass: read the whole file and count the parts, so problems are
	// reported on the page before any of the zip is sent.
	st, err := pcap.Split(file, limit, func(part int) (io.Writer, error) {
		if part > maxSplitParts {
			return nil, fmt.Errorf("this would create more than %d files; use a larger limit", maxSplitParts)
		}
		return io.Discard, nil
	})
	if err != nil {
		fail(http.StatusBadRequest, hdr.Filename+": "+err.Error())
		return
	}
	if _, err := file.Seek(0, io.SeekStart); err != nil {
		fail(http.StatusInternalServerError, err.Error())
		return
	}

	name := safeName(hdr.Filename, "capture.pcap")
	base := strings.TrimSuffix(name, path.Ext(name))
	digits := max(3, len(strconv.Itoa(st.Parts)))

	w.Header().Set("Content-Type", "application/zip")
	w.Header().Set("Content-Disposition", fmt.Sprintf("attachment; filename=%q", base+"_split.zip"))
	zw := zip.NewWriter(w)
	zw.RegisterCompressor(zip.Deflate, func(out io.Writer) (io.WriteCloser, error) {
		return flate.NewWriter(out, flate.BestSpeed)
	})
	if st.Truncated {
		zw.SetComment("The input file ended with a truncated packet record, which was dropped.")
	}
	now := time.Now()
	_, err = pcap.Split(file, limit, func(part int) (io.Writer, error) {
		return zw.CreateHeader(&zip.FileHeader{
			Name:     fmt.Sprintf("%s_%0*d.pcap", base, digits, part),
			Method:   zip.Deflate,
			Modified: now,
		})
	})
	if err == nil {
		err = zw.Close()
	}
	if err != nil {
		// The response has already started, so the error cannot be shown on the page.
		log.Printf("split %q: %v", hdr.Filename, err)
	}
}

func parseSplitLimit(f splitForm) (pcap.SplitLimit, error) {
	switch f.Mode {
	case "packets":
		n, err := strconv.Atoi(strings.TrimSpace(f.Packets))
		if err != nil || n < 1 {
			return pcap.SplitLimit{}, errors.New("packets per file must be a whole number of at least 1")
		}
		return pcap.SplitLimit{Packets: n}, nil

	case "size":
		v, err := strconv.ParseFloat(strings.TrimSpace(f.Size), 64)
		if err != nil || math.IsNaN(v) || math.IsInf(v, 0) || v <= 0 {
			return pcap.SplitLimit{}, errors.New("size per file must be a positive number")
		}
		for _, u := range sizeUnits {
			if u.Name == f.Unit {
				b := v * u.Bytes
				if b < 1000 {
					return pcap.SplitLimit{}, errors.New("size per file must be at least 1 KB")
				}
				if b > math.MaxInt64/2 {
					return pcap.SplitLimit{}, errors.New("size per file is too large")
				}
				return pcap.SplitLimit{Bytes: int64(b)}, nil
			}
		}
		return pcap.SplitLimit{}, fmt.Errorf("unknown size unit %q", f.Unit)
	}
	return pcap.SplitLimit{}, errors.New("choose whether to split by packet count or by file size")
}