warc/reader.go

   1 /*
   2 tofuproxy -- flexible HTTP/WARC proxy with TLS certificates management
   3 Copyright (C) 2021 Sergey Matveev <stargrave@stargrave.org>
   4
   5 This program is free software: you can redistribute it and/or modify
   6 it under the terms of the GNU General Public License as published by
   7 the Free Software Foundation, version 3 of the License.
   8
   9 This program is distributed in the hope that it will be useful,
  10 but WITHOUT ANY WARRANTY; without even the implied warranty of
  11 MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  12 GNU General Public License for more details.
  13
  14 You should have received a copy of the GNU General Public License
  15 along with this program.  If not, see <http://www.gnu.org/licenses/>.
  16 */
  17
  18 package warc
  19
  20 import (
  21         "bufio"
  22         "fmt"
  23         "io"
  24         "strconv"
  25         "strings"
  26 )
  27
  28 const CRLF = "\r\n"
  29
  30 type Reader struct {
  31         Path    string
  32         rrr     RawRecordReader
  33         br      *bufio.Reader
  34         offset  int64
  35         prevRec *Record
  36         offsets []Offset
  37 }
  38
  39 func NewReader(warcPath string) (*Reader, error) {
  40         rrr, err := Open(warcPath, nil, 0)
  41         if err != nil {
  42                 return nil, err
  43         }
  44         return &Reader{
  45                 Path: warcPath,
  46                 rrr:  rrr,
  47                 br:   bufio.NewReader(rrr),
  48         }, nil
  49 }
  50
  51 func (r *Reader) next() error {
  52         if r.prevRec == nil {
  53                 return nil
  54         }
  55         if _, err := r.br.Discard(int(r.prevRec.Size)); err != nil {
  56                 return err
  57         }
  58         r.offset += int64(r.prevRec.HdrLen) + r.prevRec.Size
  59         for i := 0; i < 2; i++ {
  60                 line, err := r.br.ReadString('\n')
  61                 if err != nil {
  62                         return err
  63                 }
  64                 r.offset += int64(len(line))
  65                 if line != CRLF {
  66                         return fmt.Errorf("non-CRLF: %q", line)
  67                 }
  68         }
  69         return nil
  70 }
  71
  72 func (r *Reader) ReadRecord() (*Record, io.Reader, error) {
  73         r.next()
  74         line, err := r.br.ReadString('\n')
  75         if err != nil {
  76                 return nil, nil, err
  77         }
  78         if !strings.HasPrefix(line, "WARC/") {
  79                 return nil, nil, fmt.Errorf("non-WARC header: %q", line)
  80         }
  81         hdrLines := []string{line}
  82         hdrLen := len(line)
  83         hdr := NewHeader()
  84         for {
  85                 line, err := r.br.ReadString('\n')
  86                 if err != nil {
  87                         return nil, nil, err
  88                 }
  89                 hdrLen += len(line)
  90                 if line == CRLF {
  91                         break
  92                 }
  93                 hdrLines = append(hdrLines, line)
  94                 hdr.AddLine(line)
  95         }
  96         size, err := strconv.ParseUint(hdr.Get("Content-Length"), 10, 64)
  97         if err != nil {
  98                 return nil, nil, err
  99         }
 100         rec := &Record{
 101                 WARCPath: r.Path,
 102                 Offset:   r.offset,
 103                 Size:     int64(size),
 104                 Hdr:      hdr,
 105                 HdrLen:   hdrLen,
 106                 HdrLines: hdrLines,
 107         }
 108         r.prevRec = rec
 109         return rec, &io.LimitedReader{R: r.br, N: int64(size)}, nil
 110 }
 111
 112 func (r *Reader) RecordWasRead() {
 113         r.prevRec.HdrLen = 0
 114         r.prevRec.Size = 0
 115 }
 116
 117 func (r *Reader) Close() error {
 118         err := r.rrr.Close()
 119         r.offsets = r.rrr.Offsets()
 120         return err
 121 }