path: root/image/webp/client.go

package webp

import (
	"bytes"
	"context"
	"fmt"
	"image"
	"io"
	"net"
	"net/http"
	"os"
	"os/exec"
	"path"
	"strconv"
	"syscall"
	"time"

	"code.crute.us/mcrute/golib/log"
	"golang.org/x/image/tiff"
)

// WebPConvertClient starts a webp conversion server and manages
// its lifecycle. Call Convert to convert an image.Image to a webp
// bytestream. This struct is only threadsafe after Start has been
// called.
type WebPConvertClient struct {
	p      *exec.Cmd
	c      *http.Client
	Logger log.LeveledLogger
}

// Start starts the server, cancelling the passed context will stop
// the server process. It is important to pass a real context into the
// process and not just context.Background. This function will return an
// error if setup fails but does not block.
func (c *WebPConvertClient) Start(ctx context.Context) error {
	if c.Logger == nil {
		c.Logger = &log.NoopLeveledLogger{}
	}

	sockDir, err := os.MkdirTemp("", "webp-*")
	if err != nil {
		return err
	}

	sockPath := path.Join(sockDir, "webp.sock")
	c.Logger.Debugf("Webp conversion server socket: %s", sockPath)

	c.c = &http.Client{
		Transport: &http.Transport{
			DialContext: func(ctx context.Context, _, _ string) (net.Conn, error) {
				var d net.Dialer
				return d.DialContext(ctx, "unix", sockPath)
			},
		},
	}

	c.p = exec.CommandContext(ctx, "webp-convert-server", sockPath)
	if c.p.Err != nil {
		return c.p.Err
	}

	c.p.Cancel = func() error {
		c.Logger.Debugf("Shutting down webp conversion server")
		if err := c.p.Process.Signal(syscall.SIGTERM); err != nil {
			return err
		}
		return os.ErrProcessDone
	}

	ec := make(chan error, 1)
	go func() {
		ec <- c.p.Run()
		os.RemoveAll(sockDir)
	}()

	if err := c.awaitServerReadiness(ctx, 10); err != nil {
		c.p.Cancel()
		return err
	}

	select {
	case err = <-ec:
		return err
	default:
		return nil
	}
}

// awaitServerReadiness waits for the server to become ready by pinging
// its endpoint. This should happen rather quickly, but without this
// check it can be a race to see if the server starts before the first
// call is made in some programs.
func (c *WebPConvertClient) awaitServerReadiness(ctx context.Context, maxTries int) error {
	for {
		c.Logger.Debugf("Waiting for webp conversion server to be ready")

		maxTries--
		if maxTries == 0 {
			break
		}

		if err := c.ping(ctx); err == nil {
			return nil
		}

		time.Sleep(time.Second)
	}
	return fmt.Errorf("Webp server failed to become ready")
}

func (c *WebPConvertClient) ping(ctx context.Context) error {
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, "http://unix/ping", nil)
	if err != nil {
		return err
	}

	res, err := c.c.Do(req)
	if err != nil {
		return err
	}
	defer res.Body.Close()

	if res.StatusCode != http.StatusOK {
		return fmt.Errorf("server not ready")
	}

	return nil
}

// Convert encodes an image as TIFF and sends it to the conversion
// server. It returns either an error or the webp encoded bytes from the
// input image.
func (c *WebPConvertClient) Convert(ctx context.Context, img image.Image, quality int, lossless, exact bool) ([]byte, error) {
	if c.c == nil {
		return nil, fmt.Errorf("WebPConvertClient has not been started")
	}

	bb := &bytes.Buffer{}
	if err := tiff.Encode(bb, img, &tiff.Options{Compression: tiff.Uncompressed}); err != nil {
		return nil, err
	}

	return c.ConvertBytes(ctx, bb.Bytes(), quality, lossless, exact)
}

// ConvertBytes takes an image in any format and sends it to the
// conversion server. It returns either an error or the webp encoded
// bytes from the input image.
func (c *WebPConvertClient) ConvertBytes(ctx context.Context, img []byte, quality int, lossless, exact bool) ([]byte, error) {
	if c.c == nil {
		return nil, fmt.Errorf("WebPConvertClient has not been started")
	}

	req, err := http.NewRequestWithContext(ctx, http.MethodPost, "http://unix/", bytes.NewReader(img))
	if err != nil {
		return nil, err
	}
	req.Header.Add("quality", strconv.Itoa(quality))
	req.Header.Add("lossless", fmt.Sprintf("%t", lossless))
	req.Header.Add("exact", fmt.Sprintf("%t", exact))

	res, err := c.c.Do(req)
	if err != nil {
		return nil, err
	}
	defer res.Body.Close()

	if res.StatusCode != http.StatusOK {
		e, err := io.ReadAll(res.Body)
		if err != nil {
			return nil, err
		}
		return nil, fmt.Errorf("WebPConvertClient server error: %s", e)
	}

	webp, err := io.ReadAll(res.Body)
	if err != nil {
		return nil, fmt.Errorf("WebPConvertClient error reading response: %w", err)
	}

	return webp, nil
}

// ConvertBytesStreaming takes an image in any format and sends it to
// the conversion server. It returns either an error or the webp encoded
// bytes as a ReadCloser from the input image. The caller must Close the
// returned ReadCloser or it will leak.
func (c *WebPConvertClient) ConvertBytesStreaming(ctx context.Context, img []byte, quality int, lossless, exact bool) (io.ReadCloser, error) {
	if c.c == nil {
		return nil, fmt.Errorf("WebPConvertClient has not been started")
	}

	req, err := http.NewRequestWithContext(ctx, http.MethodPost, "http://unix/", bytes.NewReader(img))
	if err != nil {
		return nil, err
	}
	req.Header.Add("quality", strconv.Itoa(quality))
	req.Header.Add("lossless", fmt.Sprintf("%t", lossless))
	req.Header.Add("exact", fmt.Sprintf("%t", exact))

	res, err := c.c.Do(req)
	if err != nil {
		return nil, err
	}

	if res.StatusCode != http.StatusOK {
		e, err := io.ReadAll(res.Body)
		if err != nil {
			return nil, err
		}
		return nil, fmt.Errorf("WebPConvertClient server error: %s", e)
	}

	return res.Body, nil
}