warc/compressed.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         "compress/gzip"
  22         "fmt"
  23         "io"
  24         "os"
  25         "path"
  26
  27         "github.com/klauspost/compress/zstd"
  28 )
  29
  30 type Compressed struct {
  31         r      io.ReadCloser
  32         fd     *os.File
  33         offset int64
  34 }
  35
  36 func (c *Compressed) Read(p []byte) (int, error) {
  37         n, err := c.r.Read(p)
  38         c.offset += int64(n)
  39         return n, err
  40 }
  41
  42 func (c *Compressed) Close() error {
  43         c.r.Close()
  44         return c.fd.Close()
  45 }
  46
  47 func (c *Compressed) Seek(offset int64, whence int) (int64, error) {
  48         if whence != io.SeekStart {
  49                 panic("can only seek from the start")
  50         }
  51         if _, err := io.CopyN(io.Discard, c, offset-c.offset); err != nil {
  52                 return 0, err
  53         }
  54         c.offset = offset
  55         return c.offset, nil
  56 }
  57
  58 func Open(warcPath string) (io.ReadSeekCloser, error) {
  59         ext := path.Ext(warcPath)
  60         switch ext {
  61         case ".warc":
  62                 return os.Open(warcPath)
  63         case ".gz":
  64                 fd, err := os.Open(warcPath)
  65                 if err != nil {
  66                         return nil, err
  67                 }
  68                 gzr, err := gzip.NewReader(fd)
  69                 if err != nil {
  70                         return nil, err
  71                 }
  72                 gzr.Multistream(true)
  73                 return &Compressed{r: gzr, fd: fd}, nil
  74         case ".zst":
  75                 fd, err := os.Open(warcPath)
  76                 if err != nil {
  77                         return nil, err
  78                 }
  79                 zstdr, err := zstd.NewReader(fd)
  80                 if err != nil {
  81                         return nil, err
  82                 }
  83                 return &Compressed{r: zstdr.IOReadCloser(), fd: fd}, nil
  84         }
  85         return nil, fmt.Errorf("unknown extensions: %s", ext)
  86 }