warc/gzip.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         "bufio"
  21         "compress/gzip"
  22         "io"
  23         "os"
  24 )
  25
  26 type CountableReader struct {
  27         *bufio.Reader
  28         n int64
  29 }
  30
  31 func (r *CountableReader) ReadByte() (b byte, err error) {
  32         b, err = r.Reader.ReadByte()
  33         if err == nil {
  34                 r.n++
  35         }
  36         return b, err
  37 }
  38
  39 func (r *CountableReader) Read(p []byte) (n int, err error) {
  40         n, err = r.Reader.Read(p)
  41         r.n += int64(n)
  42         return
  43 }
  44
  45 type GZIPReader struct {
  46         fd      *os.File
  47         r       io.Reader
  48         offsets []Offset
  49 }
  50
  51 func (r *GZIPReader) Read(p []byte) (n int, err error) {
  52         return r.r.Read(p)
  53 }
  54
  55 func (r *GZIPReader) Close() error {
  56         return r.fd.Close()
  57 }
  58
  59 func (r *GZIPReader) Offsets() []Offset {
  60         return r.offsets
  61 }
  62
  63 func NewGZIPReader(
  64         warcPath string,
  65         offsets []Offset,
  66         uOffset int64,
  67 ) (*GZIPReader, error) {
  68         var offZ, offU int64
  69         for _, off := range offsets {
  70                 if uOffset < offU+off.U {
  71                         break
  72                 }
  73                 offU += off.U
  74                 offZ += off.Z
  75         }
  76         fd, err := os.Open(warcPath)
  77         if err != nil {
  78                 return nil, err
  79         }
  80         if _, err = fd.Seek(offZ, io.SeekStart); err != nil {
  81                 fd.Close()
  82                 return nil, err
  83         }
  84         cr := &CountableReader{Reader: bufio.NewReader(fd)}
  85         z, err := gzip.NewReader(cr)
  86         if err != nil {
  87                 fd.Close()
  88                 return nil, err
  89         }
  90         r, w := io.Pipe()
  91         gr := GZIPReader{r: r}
  92         go func() {
  93                 z.Multistream(false)
  94                 var offset, offsetPrev int64
  95                 for {
  96                         written, err := io.Copy(w, z)
  97                         if err != nil {
  98                                 w.CloseWithError(err)
  99                                 return
 100                         }
 101                         offset = cr.n
 102                         gr.offsets = append(gr.offsets, Offset{offset - offsetPrev, written})
 103                         offsetPrev = offset
 104                         err = z.Reset(cr)
 105                         if err != nil {
 106                                 if err == io.EOF {
 107                                         break
 108                                 }
 109                                 w.CloseWithError(err)
 110                                 return
 111                         }
 112                         z.Multistream(false)
 113                 }
 114                 w.CloseWithError(io.EOF)
 115         }()
 116         _, err = io.CopyN(io.Discard, r, uOffset-offU)
 117         if err != nil {
 118                 fd.Close()
 119                 return nil, err
 120         }
 121         return &gr, nil
 122 }