DankMaterialShell/core/internal/netfetch/netfetch.go

package netfetch

import (
	"context"
	"fmt"
	"io"
	"net"
	"net/http"
	"os"
	"path/filepath"
	"time"
)

const DefaultUserAgent = "DankMaterialShell/1.0 (Linux)"

type Options struct {
	Headers        map[string]string
	UserAgent      string
	Timeout        time.Duration
	ConnectTimeout time.Duration
	IPv4Only       bool
}

func (o Options) client() *http.Client {
	connect := o.ConnectTimeout
	if connect <= 0 {
		connect = 5 * time.Second
	}

	dialer := &net.Dialer{Timeout: connect}
	transport := &http.Transport{DialContext: dialer.DialContext}
	if o.IPv4Only {
		transport.DialContext = func(ctx context.Context, network, addr string) (net.Conn, error) {
			return dialer.DialContext(ctx, "tcp4", addr)
		}
	}
	return &http.Client{Transport: transport}
}

func (o Options) request(ctx context.Context, url string) (*http.Request, error) {
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
	if err != nil {
		return nil, fmt.Errorf("invalid request: %w", err)
	}

	agent := o.UserAgent
	if agent == "" {
		agent = DefaultUserAgent
	}
	req.Header.Set("User-Agent", agent)

	for name, value := range o.Headers {
		req.Header.Set(name, value)
	}
	return req, nil
}

// Closing the body is what cancels the timeout context.
func Open(ctx context.Context, url string, opts Options) (io.ReadCloser, error) {
	if opts.Timeout <= 0 {
		return open(ctx, url, opts)
	}

	ctx, cancel := context.WithTimeout(ctx, opts.Timeout)
	body, err := open(ctx, url, opts)
	if err != nil {
		cancel()
		return nil, err
	}
	return &cancelOnClose{ReadCloser: body, cancel: cancel}, nil
}

type cancelOnClose struct {
	io.ReadCloser
	cancel context.CancelFunc
}

func (c *cancelOnClose) Close() error {
	err := c.ReadCloser.Close()
	c.cancel()
	return err
}

func open(ctx context.Context, url string, opts Options) (io.ReadCloser, error) {
	req, err := opts.request(ctx, url)
	if err != nil {
		return nil, err
	}

	resp, err := opts.client().Do(req)
	if err != nil {
		return nil, fmt.Errorf("download failed: %w", err)
	}

	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
		resp.Body.Close()
		return nil, fmt.Errorf("HTTP %d", resp.StatusCode)
	}
	return resp.Body, nil
}

func Bytes(ctx context.Context, url string, opts Options) ([]byte, error) {
	body, err := Open(ctx, url, opts)
	if err != nil {
		return nil, err
	}
	defer body.Close()
	return io.ReadAll(body)
}

func ToWriter(ctx context.Context, url string, opts Options, w io.Writer) error {
	body, err := Open(ctx, url, opts)
	if err != nil {
		return err
	}
	defer body.Close()

	_, err = io.Copy(w, body)
	return err
}

func ToFile(ctx context.Context, url string, opts Options, path string) error {
	if dir := filepath.Dir(path); dir != "." {
		if err := os.MkdirAll(dir, 0o755); err != nil {
			return fmt.Errorf("mkdir failed: %w", err)
		}
	}

	f, err := os.Create(path)
	if err != nil {
		return fmt.Errorf("create failed: %w", err)
	}
	defer f.Close()

	if err := ToWriter(ctx, url, opts, f); err != nil {
		os.Remove(path)
		return err
	}
	return nil
}