compressed.go

   1 /*
   2 tofuproxy -- flexible HTTP/HTTPS proxy, TLS terminator, X.509 TOFU
   3              manager, WARC/geminispace browser
   4 Copyright (C) 2021-2023 Sergey Matveev <stargrave@stargrave.org>
   5
   6 This program is free software: you can redistribute it and/or modify
   7 it under the terms of the GNU General Public License as published by
   8 the Free Software Foundation, version 3 of the License.
   9
  10 This program is distributed in the hope that it will be useful,
  11 but WITHOUT ANY WARRANTY; without even the implied warranty of
  12 MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  13 GNU General Public License for more details.
  14
  15 You should have received a copy of the GNU General Public License
  16 along with this program.  If not, see <http://www.gnu.org/licenses/>.
  17 */
  18
  19 package warc
  20
  21 import (
  22         "bufio"
  23         "bytes"
  24         "io"
  25         "log"
  26         "os"
  27         "os/exec"
  28         "strconv"
  29         "strings"
  30         "sync"
  31 )
  32
  33 type CompressedReader struct {
  34         cmd     *exec.Cmd
  35         fd      *os.File
  36         stdout  io.ReadCloser
  37         offsets []Offset
  38
  39         offW      *os.File
  40         offReader sync.WaitGroup
  41 }
  42
  43 func NewCompressedReader(
  44         warcPath, unCmd string,
  45         offsets []Offset,
  46         uOffset int64,
  47 ) (*CompressedReader, error) {
  48         var offZ, offU int64
  49         for _, off := range offsets {
  50                 if uOffset < offU+off.U {
  51                         break
  52                 }
  53                 offU += off.U
  54                 offZ += off.Z
  55         }
  56         fd, err := os.Open(warcPath)
  57         if err != nil {
  58                 return nil, err
  59         }
  60         var dict []byte
  61         if len(offsets) > 0 && offsets[0].U == 0 {
  62                 dict = make([]byte, offsets[0].Z)
  63                 if _, err = io.ReadFull(fd, dict); err != nil {
  64                         fd.Close()
  65                         return nil, err
  66                 }
  67         }
  68         if _, err = fd.Seek(offZ, io.SeekStart); err != nil {
  69                 fd.Close()
  70                 return nil, err
  71         }
  72         cmd := exec.Command(unCmd)
  73         stdout, err := cmd.StdoutPipe()
  74         if err != nil {
  75                 fd.Close()
  76                 return nil, err
  77         }
  78         if dict == nil {
  79                 cmd.Stdin = fd
  80         } else {
  81                 cmd.Stdin = io.MultiReader(bytes.NewReader(dict), fd)
  82         }
  83         if offsets == nil {
  84                 offR, offW, err := os.Pipe()
  85                 if err != nil {
  86                         fd.Close()
  87                         return nil, err
  88                 }
  89                 cmd.ExtraFiles = append(cmd.ExtraFiles, offW)
  90                 err = cmd.Start()
  91                 if err != nil {
  92                         fd.Close()
  93                         offW.Close()
  94                         return nil, err
  95                 }
  96                 r := CompressedReader{
  97                         cmd:    cmd,
  98                         fd:     fd,
  99                         stdout: stdout,
 100                         offW:   offW,
 101                 }
 102                 r.offReader.Add(1)
 103                 go r.offsetsReader(offR)
 104                 return &r, nil
 105         }
 106         err = cmd.Start()
 107         if err != nil {
 108                 fd.Close()
 109                 return nil, err
 110         }
 111         _, err = io.CopyN(io.Discard, stdout, uOffset-offU)
 112         if err != nil {
 113                 cmd.Process.Kill()
 114                 fd.Close()
 115                 return nil, err
 116         }
 117         return &CompressedReader{cmd: cmd, fd: fd, stdout: stdout}, nil
 118 }
 119
 120 func (r *CompressedReader) offsetsReader(offsets *os.File) {
 121         scanner := bufio.NewScanner(offsets)
 122         for scanner.Scan() {
 123                 l := scanner.Text()
 124                 cols := strings.Split(l, "\t")
 125                 if len(cols) != 2 {
 126                         log.Println("len(cols) != 2:", l)
 127                         continue
 128                 }
 129                 z, err := strconv.ParseUint(cols[0], 10, 64)
 130                 if err != nil {
 131                         log.Println(err)
 132                         continue
 133                 }
 134                 u, err := strconv.ParseUint(cols[1], 10, 64)
 135                 if err != nil {
 136                         log.Println(err)
 137                         continue
 138                 }
 139                 r.offsets = append(r.offsets, Offset{int64(z), int64(u)})
 140         }
 141         err := scanner.Err()
 142         if err != nil {
 143                 log.Println(err)
 144         }
 145         r.offReader.Done()
 146 }
 147
 148 func (r *CompressedReader) Read(p []byte) (int, error) {
 149         return r.stdout.Read(p)
 150 }
 151
 152 func (r *CompressedReader) Close() error {
 153         err := r.cmd.Process.Kill()
 154         r.stdout.Close()
 155         r.fd.Close()
 156         r.offW.Close()
 157         r.offReader.Wait()
 158         return err
 159 }
 160
 161 func (r *CompressedReader) Offsets() []Offset {
 162         return r.offsets
 163 }