warc.go 2.22 KB
Newer Older
ale's avatar
ale committed
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96
package crawl

import (
	"fmt"
	"io"
	"time"

	"compress/gzip"

	"code.google.com/p/go-uuid/uuid"
)

var (
	warcTimeFmt      = time.RFC3339
	warcVersion      = "WARC/1.0"
	warcContentTypes = map[string]string{
		"warcinfo": "application/warc-fields",
		"response": "application/http; msgtype=response",
		"request":  "application/http; msgtype=request",
		"metadata": "application/warc-fields",
	}
)

// A Warc header. Header field names are case-sensitive.
type WarcHeader map[string]string

// Set a header to the specified value. Multiple values are not
// supported.
func (h WarcHeader) Set(key, value string) {
	h[key] = value

	// Keep Content-Type in sync with WARC-Type.
	if key == "WARC-Type" {
		if ct, ok := warcContentTypes[value]; ok {
			h["Content-Type"] = ct
		} else {
			h["Content-Type"] = "application/octet-stream"
		}
	}
}

// Get the value of a header. If not found, returns an empty string.
func (h WarcHeader) Get(key string) string {
	return h[key]
}

// Encode the header to a Writer.
func (h WarcHeader) Encode(w io.Writer) {
	fmt.Fprintf(w, "%s\r\n", warcVersion)
	for hdr, value := range h {
		fmt.Fprintf(w, "%s: %s\r\n", hdr, value)
	}
	fmt.Fprintf(w, "\r\n")
}

// NewWarcHeader returns a WarcHeader with its own unique ID and the
// current timestamp.
func NewWarcHeader() WarcHeader {
	h := make(WarcHeader)
	h.Set("WARC-Record-ID", fmt.Sprintf("<%s>", uuid.NewUUID().URN()))
	h.Set("WARC-Date", time.Now().Format(warcTimeFmt))
	h.Set("Content-Type", "application/octet-stream")
	return h
}

// WarcWriter can write records to a file in WARC format.
type WarcWriter struct {
	writer io.WriteCloser
}

type recordWriter struct {
	io.Writer
}

func (rw *recordWriter) Close() error {
	// Add the end-of-record marker.
	fmt.Fprintf(rw, "\r\n\r\n")
	return nil
}

// NewRecord starts a new WARC record with the provided header. The
// caller must call Close on the returned writer before creating the
// next record.
func (w *WarcWriter) NewRecord(hdr WarcHeader) io.WriteCloser {
	hdr.Encode(w.writer)
	return &recordWriter{w.writer}
}

// Close the WARC writer and flush all buffers.
func (w *WarcWriter) Close() error {
	return w.writer.Close()
}

func NewWarcWriter(w io.WriteCloser) *WarcWriter {
	return &WarcWriter{gzip.NewWriter(w)}
}