warc/record.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         "io"
  22         "strings"
  23 )
  24
  25 type Record struct {
  26         WARCPath string
  27         Offset   int64
  28         Size     int64
  29
  30         Hdr      Header
  31         HdrLen   int
  32         HdrLines []string
  33
  34         Continuations []*Record
  35 }
  36
  37 func (rec *Record) URI() string {
  38         return strings.Trim(rec.Hdr.Get("WARC-Target-URI"), "<>")
  39 }
  40
  41 func (rec *Record) TotalSize() int64 {
  42         s := rec.Size
  43         for _, r := range rec.Continuations {
  44                 s += r.Size
  45         }
  46         return s
  47 }
  48
  49 type SelfRecordReader struct {
  50         lr  *io.LimitedReader
  51         rrr io.ReadCloser
  52 }
  53
  54 func (srr *SelfRecordReader) Read(p []byte) (n int, err error) {
  55         n, err = srr.rrr.Read(p)
  56         if err != nil {
  57                 srr.Close()
  58         }
  59         return
  60 }
  61
  62 func (srr *SelfRecordReader) Close() error {
  63         return srr.rrr.Close()
  64 }
  65
  66 func (rec *Record) selfReader(noHdr bool, offsets []Offset) (*SelfRecordReader, error) {
  67         offset := rec.Offset
  68         if noHdr {
  69                 offset += int64(rec.HdrLen)
  70         }
  71         rrr, err := Open(rec.WARCPath, offsets, offset)
  72         if err != nil {
  73                 return nil, err
  74         }
  75         return &SelfRecordReader{lr: &io.LimitedReader{R: rrr, N: rec.Size}, rrr: rrr}, nil
  76 }
  77
  78 type RecordReader struct {
  79         r    io.Reader
  80         srrs []*SelfRecordReader
  81 }
  82
  83 func (rec *Record) Reader(
  84         noHdr bool,
  85         warcOffsets map[string][]Offset,
  86 ) (*RecordReader, error) {
  87         srrs := make([]*SelfRecordReader, 0, 1+len(rec.Continuations))
  88         rs := make([]io.Reader, 0, 1+len(rec.Continuations))
  89         for i, r := range append([]*Record{rec}, rec.Continuations...) {
  90                 if i > 0 {
  91                         noHdr = true
  92                 }
  93                 srr, err := r.selfReader(noHdr, warcOffsets[rec.WARCPath])
  94                 if err != nil {
  95                         for _, srr := range srrs {
  96                                 srr.Close()
  97                         }
  98                         return nil, err
  99                 }
 100                 srrs = append(srrs, srr)
 101                 rs = append(rs, srr)
 102         }
 103         return &RecordReader{r: io.MultiReader(rs...), srrs: srrs}, nil
 104 }
 105
 106 func (rr *RecordReader) Read(p []byte) (n int, err error) {
 107         n, err = rr.r.Read(p)
 108         if err != nil {
 109                 rr.Close()
 110         }
 111         return
 112 }
 113
 114 func (rr *RecordReader) Close() error {
 115         for _, srr := range rr.srrs {
 116                 srr.Close()
 117         }
 118         return nil
 119 }