warc/record.go

   1 /*
   2 tofuproxy -- flexible HTTP/HTTPS proxy, TLS terminator, X.509 TOFU
   3              manager, WARC/geminispace browser
   4 Copyright (C) 2021-2022 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.rrr.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         return &SelfRecordReader{lr: &io.LimitedReader{R: rrr, N: rec.Size}, rrr: rrr}, nil
  77 }
  78
  79 type RecordReader struct {
  80         r    io.Reader
  81         srrs []*SelfRecordReader
  82 }
  83
  84 func (rec *Record) Reader(
  85         noHdr bool,
  86         warcOffsets map[string][]Offset,
  87 ) (*RecordReader, error) {
  88         srrs := make([]*SelfRecordReader, 0, 1+len(rec.Continuations))
  89         rs := make([]io.Reader, 0, 1+len(rec.Continuations))
  90         for i, r := range append([]*Record{rec}, rec.Continuations...) {
  91                 if i > 0 {
  92                         noHdr = true
  93                 }
  94                 srr, err := r.selfReader(noHdr, warcOffsets[rec.WARCPath])
  95                 if err != nil {
  96                         for _, srr := range srrs {
  97                                 srr.Close()
  98                         }
  99                         return nil, err
 100                 }
 101                 srrs = append(srrs, srr)
 102                 rs = append(rs, srr)
 103         }
 104         return &RecordReader{r: io.MultiReader(rs...), srrs: srrs}, nil
 105 }
 106
 107 func (rr *RecordReader) Read(p []byte) (n int, err error) {
 108         n, err = rr.r.Read(p)
 109         if err != nil {
 110                 rr.Close()
 111         }
 112         return
 113 }
 114
 115 func (rr *RecordReader) Close() error {
 116         for _, srr := range rr.srrs {
 117                 srr.Close()
 118         }
 119         return nil
 120 }