write.go

package warc

import (
	"bufio"
	"bytes"
	"fmt"
	"io"
	"strconv"
	"strings"
	"time"

	"github.com/google/uuid"
	"github.com/saveweb/gowarc/pkg/spooledtempfile"
)

type Compressor interface {
	io.Writer // Embedded so the compressor can be used as an io.Writer and its underlying writer can be replaced via Reset.
	io.Closer
	Reset(io.Writer)
}

const (
	warcRecordVersionLine = "WARC/1.1\r\n"
	warcHeaderEnd         = "\r\n"
	warcRecordTrailer     = "\r\n\r\n"
)

// Writer writes WARC records to WARC files.
type Writer struct {
	Compressor      Compressor
	BufWriter       *bufio.Writer
	FileName        string
	DigestAlgorithm DigestAlgorithm
	ParallelGZIP    bool
}

type FeedbackEvent []RecordEvent // WARC-Record-ID header values of the records that have been written

// RecordBatch is a structure that contains a bunch of
// records to be written at the same time, and a common
// capture timestamp. FeedbackChan is used to signal
// when the records have been written.
type RecordBatch struct {
	FeedbackChan chan FeedbackEvent
	CaptureTime  string
	Records      []*Record
}

type RecordInfo struct {
	Header  Header
	Offset  int64  // Offset of the record start (-1 if WARC file type is not supported yet)
	Size    int64  // COMPRESSED size of the record (gzip member): header + deflate data + trailer. (-1 if WARC file type is not supported yet)
	Version string // WARC/1.0, WARC/1.1 ...
}

// Record represents a WARC record.
type Record struct {
	RecordInfo
	Content spooledtempfile.ReadWriteSeekCloser
}

// a record that has been written to a WARC file.
type RecordEvent struct {
	RecordInfo
	WARCFilename string // WARC filename, no .open extension
}

// WriteRecord writes a record to the underlying WARC file and flushes the data.
// A record consists of a version string, the record header followed by a
// record content block and two newlines:
//
//	Version CLRF
//	Header-Key: Header-Value CLRF
//	CLRF
//	Content
//	CLRF
//	CLRF
func (w *Writer) WriteRecord(r *Record) (recordID string, err error) {
	defer r.Content.Close()

	// Add the mandatories headers
	if r.Header.Get("WARC-Date") == "" {
		r.Header.Set("WARC-Date", time.Now().UTC().Format(time.RFC3339Nano))
	}

	if r.Header.Get("WARC-Type") == "" {
		r.Header.Set("WARC-Type", "resource")
	}

	if r.Header.Get("WARC-Record-ID") == "" {
		recordID = uuid.NewString()
		r.Header.Set("WARC-Record-ID", "<urn:uuid:"+recordID+">")
	}

	var contentLength int64
	if CL := r.Header.Get("Content-Length"); CL != "" {
		contentLength, err = strconv.ParseInt(CL, 10, 64)
	}
	if r.Header.Get("Content-Length") == "" || err != nil {
		contentLength = getContentLength(r.Content)
		r.Header.Set("Content-Length", strconv.FormatInt(contentLength, 10))
	}

	if r.Header.Get("WARC-Block-Digest") == "" {
		r.Content.Seek(0, 0)

		digest, err := GetDigest(r.Content, w.DigestAlgorithm)
		if err != nil {
			return recordID, err
		}

		r.Header.Set("WARC-Block-Digest", digest)
	}

	headerBlock := serializedRecordHeader(r.Header)
	w.setCompressorContentSize(int64(len(headerBlock)) + contentLength + int64(len(warcRecordTrailer)))

	r.Content.Seek(0, 0)
	recordReader := io.MultiReader(
		bytes.NewReader(headerBlock),
		r.Content,
		strings.NewReader(warcRecordTrailer),
	)
	if _, err = io.Copy(w.BufWriter, recordReader); err != nil {
		return recordID, err
	}

	if contentLength > 0 {
		DataTotal.Add(contentLength)
	}

	// Flush data
	err = w.FlushAndCloseCompressor()
	if err != nil {
		return recordID, err
	}

	return recordID, nil
}

func serializedRecordHeader(header Header) []byte {
	var buf bytes.Buffer
	buf.WriteString(warcRecordVersionLine)
	for key, values := range header {
		for _, value := range values {
			buf.WriteString(key)
			buf.WriteString(": ")
			buf.WriteString(value)
			buf.WriteString("\r\n")
		}
	}
	buf.WriteString(warcHeaderEnd)
	return buf.Bytes()
}

// Must be called before any data is written to the compressor.
func (w *Writer) setCompressorContentSize(size int64) {
	if w.Compressor == nil {
		return
	}

	compressor, ok := w.Compressor.(*sizedZstdWriter)
	if !ok {
		return
	}
	compressor.SetContentSize(size)
}

// WriteInfoRecord method can be used to write an information record to the WARC file and flush the data
func (w *Writer) WriteInfoRecord(payload Header) (recordID string, err error) {
	// Initialize the record
	infoRecord := NewRecord("")

	// Set the headers
	infoRecord.Header.Set("WARC-Date", time.Now().UTC().Format(time.RFC3339Nano))
	infoRecord.Header.Set("WARC-Filename", strings.TrimSuffix(w.FileName, ".open"))
	infoRecord.Header.Set("WARC-Type", "warcinfo")
	infoRecord.Header.Set("Content-Type", "application/warc-fields")

	// Write the payload
	for k, vv := range payload {
		for _, v := range vv {
			fmt.Fprintf(infoRecord.Content, "%s: %s\r\n", k, v)
		}
	}

	// Generate WARC-Block-Digest
	digest, err := GetDigest(infoRecord.Content, w.DigestAlgorithm)
	if err != nil {
		return recordID, err
	}

	infoRecord.Header.Set("WARC-Block-Digest", digest)

	// Finally, write the record and flush the data
	recordID, err = w.WriteRecord(infoRecord)
	if err != nil {
		return recordID, err
	}

	return recordID, err
}