warc/compressed.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         "bytes"
  22         "io"
  23         "log"
  24         "os"
  25         "os/exec"
  26         "strconv"
  27         "strings"
  28         "sync"
  29 )
  30
  31 type CompressedReader struct {
  32         cmd     *exec.Cmd
  33         fd      *os.File
  34         stdout  io.ReadCloser
  35         offsets []Offset
  36
  37         offW      *os.File
  38         offReader sync.WaitGroup
  39 }
  40
  41 func NewCompressedReader(
  42         warcPath, unCmd string,
  43         offsets []Offset,
  44         uOffset int64,
  45 ) (*CompressedReader, error) {
  46         var offZ, offU int64
  47         for _, off := range offsets {
  48                 if uOffset < offU+off.U {
  49                         break
  50                 }
  51                 offU += off.U
  52                 offZ += off.Z
  53         }
  54         fd, err := os.Open(warcPath)
  55         if err != nil {
  56                 return nil, err
  57         }
  58         var dict []byte
  59         if len(offsets) > 0 && offsets[0].U == 0 {
  60                 dict = make([]byte, offsets[0].Z)
  61                 if _, err = io.ReadFull(fd, dict); err != nil {
  62                         fd.Close()
  63                         return nil, err
  64                 }
  65         }
  66         if _, err = fd.Seek(offZ, io.SeekStart); err != nil {
  67                 fd.Close()
  68                 return nil, err
  69         }
  70         cmd := exec.Command(unCmd)
  71         stdout, err := cmd.StdoutPipe()
  72         if err != nil {
  73                 fd.Close()
  74                 return nil, err
  75         }
  76         if dict == nil {
  77                 cmd.Stdin = fd
  78         } else {
  79                 cmd.Stdin = io.MultiReader(bytes.NewReader(dict), fd)
  80         }
  81         if offsets == nil {
  82                 offR, offW, err := os.Pipe()
  83                 if err != nil {
  84                         fd.Close()
  85                         return nil, err
  86                 }
  87                 cmd.ExtraFiles = append(cmd.ExtraFiles, offW)
  88                 err = cmd.Start()
  89                 if err != nil {
  90                         fd.Close()
  91                         offW.Close()
  92                         return nil, err
  93                 }
  94                 r := CompressedReader{
  95                         cmd:    cmd,
  96                         fd:     fd,
  97                         stdout: stdout,
  98                         offW:   offW,
  99                 }
 100                 r.offReader.Add(1)
 101                 go r.offsetsReader(offR)
 102                 return &r, nil
 103         }
 104         err = cmd.Start()
 105         if err != nil {
 106                 fd.Close()
 107                 return nil, err
 108         }
 109         _, err = io.CopyN(io.Discard, stdout, uOffset-offU)
 110         if err != nil {
 111                 cmd.Process.Kill()
 112                 fd.Close()
 113                 return nil, err
 114         }
 115         return &CompressedReader{cmd: cmd, fd: fd, stdout: stdout}, nil
 116 }
 117
 118 func (r *CompressedReader) offsetsReader(offsets *os.File) {
 119         scanner := bufio.NewScanner(offsets)
 120         for scanner.Scan() {
 121                 l := scanner.Text()
 122                 cols := strings.Split(l, "\t")
 123                 if len(cols) != 2 {
 124                         log.Println("len(cols) != 2:", l)
 125                         continue
 126                 }
 127                 z, err := strconv.ParseUint(cols[0], 10, 64)
 128                 if err != nil {
 129                         log.Println(err)
 130                         continue
 131                 }
 132                 u, err := strconv.ParseUint(cols[1], 10, 64)
 133                 if err != nil {
 134                         log.Println(err)
 135                         continue
 136                 }
 137                 r.offsets = append(r.offsets, Offset{int64(z), int64(u)})
 138         }
 139         err := scanner.Err()
 140         if err != nil {
 141                 log.Println(err)
 142         }
 143         r.offReader.Done()
 144 }
 145
 146 func (r *CompressedReader) Read(p []byte) (int, error) {
 147         return r.stdout.Read(p)
 148 }
 149
 150 func (r *CompressedReader) Close() error {
 151         err := r.cmd.Process.Kill()
 152         r.stdout.Close()
 153         r.fd.Close()
 154         r.offW.Close()
 155         r.offReader.Wait()
 156         return err
 157 }
 158
 159 func (r *CompressedReader) Offsets() []Offset {
 160         return r.offsets
 161 }