warc/record.go

   1 /*
   2 tofuproxy -- flexible HTTP/HTTPS proxy, TLS terminator, X.509 TOFU
   3              manager, WARC/geminispace browser
   4 Copyright (C) 2021-2023 Sergey Matveev <stargrave@stargrave.org>
   5
   6 This program is free software: you can redistribute it and/or modify
   7 it under the terms of the GNU General Public License as published by
   8 the Free Software Foundation, version 3 of the License.
   9
  10 This program is distributed in the hope that it will be useful,
  11 but WITHOUT ANY WARRANTY; without even the implied warranty of
  12 MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  13 GNU General Public License for more details.
  14
  15 You should have received a copy of the GNU General Public License
  16 along with this program.  If not, see <http://www.gnu.org/licenses/>.
  17 */
  18
  19 package warc
  20
  21 import (
  22         "io"
  23         "strings"
  24 )
  25
  26 type Record struct {
  27         WARCPath string
  28         Offset   int64
  29         Size     int64
  30
  31         Hdr      Header
  32         HdrLen   int
  33         HdrLines []string
  34
  35         Continuations []*Record
  36 }
  37
  38 func (rec *Record) URI() string {
  39         return strings.Trim(rec.Hdr.Get("WARC-Target-URI"), "<>")
  40 }
  41
  42 func (rec *Record) TotalSize() int64 {
  43         s := rec.Size
  44         for _, r := range rec.Continuations {
  45                 s += r.Size
  46         }
  47         return s
  48 }
  49
  50 type SelfRecordReader struct {
  51         lr  *io.LimitedReader
  52         rrr io.ReadCloser
  53 }
  54
  55 func (srr *SelfRecordReader) Read(p []byte) (n int, err error) {
  56         n, err = srr.lr.Read(p)
  57         if err != nil {
  58                 srr.Close()
  59         }
  60         return
  61 }
  62
  63 func (srr *SelfRecordReader) Close() error {
  64         return srr.rrr.Close()
  65 }
  66
  67 func (rec *Record) selfReader(noHdr bool, offsets []Offset) (*SelfRecordReader, error) {
  68         offset := rec.Offset
  69         if noHdr {
  70                 offset += int64(rec.HdrLen)
  71         }
  72         rrr, err := Open(rec.WARCPath, offsets, offset)
  73         if err != nil {
  74                 return nil, err
  75         }
  76         size := rec.Size
  77         if !noHdr {
  78                 size += int64(rec.HdrLen)
  79         }
  80         return &SelfRecordReader{lr: &io.LimitedReader{R: rrr, N: size}, rrr: rrr}, nil
  81 }
  82
  83 type RecordReader struct {
  84         r    io.Reader
  85         srrs []*SelfRecordReader
  86 }
  87
  88 func (rec *Record) Reader(
  89         noHdr bool,
  90         warcOffsets map[string][]Offset,
  91 ) (*RecordReader, error) {
  92         srrs := make([]*SelfRecordReader, 0, 1+len(rec.Continuations))
  93         rs := make([]io.Reader, 0, 1+len(rec.Continuations))
  94         for i, r := range append([]*Record{rec}, rec.Continuations...) {
  95                 if i > 0 {
  96                         noHdr = true
  97                 }
  98                 srr, err := r.selfReader(noHdr, warcOffsets[rec.WARCPath])
  99                 if err != nil {
 100                         for _, srr := range srrs {
 101                                 srr.Close()
 102                         }
 103                         return nil, err
 104                 }
 105                 srrs = append(srrs, srr)
 106                 rs = append(rs, srr)
 107         }
 108         return &RecordReader{r: io.MultiReader(rs...), srrs: srrs}, nil
 109 }
 110
 111 func (rr *RecordReader) Read(p []byte) (n int, err error) {
 112         n, err = rr.r.Read(p)
 113         if err != nil {
 114                 rr.Close()
 115         }
 116         return
 117 }
 118
 119 func (rr *RecordReader) Close() error {
 120         for _, srr := range rr.srrs {
 121                 srr.Close()
 122         }
 123         return nil
 124 }