warc/record.go

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