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