fork of https://github.com/sourcegraph/zoekt
1package main
2
3import (
4 "archive/tar"
5 "archive/zip"
6 "bytes"
7 "compress/gzip"
8 "fmt"
9 "io"
10 "net/http"
11 "net/url"
12 "os"
13 "strings"
14)
15
16type Archive interface {
17 Next() (*File, error)
18 Close() error
19}
20
21type File struct {
22 io.ReadCloser
23 Name string
24 Size int64
25}
26
27type tarArchive struct {
28 io.Closer
29 tr *tar.Reader
30}
31
32func (a *tarArchive) Next() (*File, error) {
33 for {
34 hdr, err := a.tr.Next()
35 if err != nil {
36 return nil, err
37 }
38
39 // We only care about files
40 if hdr.Typeflag != tar.TypeReg && hdr.Typeflag != tar.TypeRegA {
41 continue
42 }
43
44 return &File{
45 ReadCloser: io.NopCloser(a.tr),
46 Name: hdr.Name,
47 Size: hdr.Size,
48 }, nil
49 }
50}
51
52type zipArchive struct {
53 io.Closer
54 files []*zip.File
55}
56
57func (a *zipArchive) Next() (*File, error) {
58 if len(a.files) == 0 {
59 return nil, io.EOF
60 }
61
62 f := a.files[0]
63 a.files = a.files[1:]
64
65 r, err := f.Open()
66 if err != nil {
67 return nil, err
68 }
69
70 return &File{
71 ReadCloser: r,
72 Name: f.Name,
73 Size: int64(f.UncompressedSize64),
74 }, nil
75}
76
77func newZipArchive(r io.Reader, closer io.Closer) (*zipArchive, error) {
78 f, ok := r.(interface {
79 io.ReaderAt
80 Stat() (os.FileInfo, error)
81 })
82 if !ok {
83 return nil, fmt.Errorf("streaming zip files not supported")
84 }
85
86 fi, err := f.Stat()
87 if err != nil {
88 return nil, err
89 }
90
91 zr, err := zip.NewReader(f, fi.Size())
92 if err != nil {
93 return nil, err
94 }
95
96 // Filter out non files
97 files := zr.File[:0]
98 for _, f := range zr.File {
99 if f.Mode().IsRegular() {
100 files = append(files, f)
101 }
102 }
103
104 return &zipArchive{
105 Closer: closer,
106 files: files,
107 }, nil
108}
109
110func detectContentType(r io.Reader) (string, io.Reader, error) {
111 var buf [512]byte
112 n, err := io.ReadFull(r, buf[:])
113 if err != nil && err != io.ErrUnexpectedEOF {
114 return "", nil, err
115 }
116
117 ct := http.DetectContentType(buf[:n])
118
119 // If we are a seeker, we can just undo our read
120 if s, ok := r.(io.Seeker); ok {
121 _, err := s.Seek(int64(-n), io.SeekCurrent)
122 return ct, r, err
123 }
124
125 // Otherwise return a new reader which merges in the read bytes
126 return ct, io.MultiReader(bytes.NewReader(buf[:n]), r), nil
127}
128
129func openReader(u string) (io.ReadCloser, error) {
130 if strings.HasPrefix(u, "https://") || strings.HasPrefix(u, "http://") {
131 resp, err := http.Get(u)
132 if err != nil {
133 return nil, err
134 }
135 if resp.StatusCode < 200 || resp.StatusCode >= 300 {
136 b, err := io.ReadAll(io.LimitReader(resp.Body, 1024))
137 _ = resp.Body.Close()
138 if err != nil {
139 return nil, err
140 }
141 return nil, &url.Error{
142 Op: "Get",
143 URL: u,
144 Err: fmt.Errorf("%s: %s", resp.Status, string(b)),
145 }
146 }
147 return resp.Body, nil
148 } else if u == "-" {
149 return io.NopCloser(os.Stdin), nil
150 }
151
152 return os.Open(u)
153}
154
155// openArchive opens the tar at the URL or filepath u. Also supported is tgz
156// files over http.
157func openArchive(u string) (ar Archive, err error) {
158 readCloser, err := openReader(u)
159 if err != nil {
160 return nil, err
161 }
162 defer func() {
163 if err != nil {
164 _ = readCloser.Close()
165 }
166 }()
167
168 ct, r, err := detectContentType(readCloser)
169 if err != nil {
170 return nil, err
171 }
172 switch ct {
173 case "application/x-gzip":
174 r, err = gzip.NewReader(r)
175 if err != nil {
176 return nil, err
177 }
178
179 case "application/zip":
180 return newZipArchive(r, readCloser)
181 }
182
183 return &tarArchive{
184 Closer: readCloser,
185 tr: tar.NewReader(r),
186 }, nil
187}