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