github.com/diamondburned/arikawa@v1.3.14/api/image.go (about)

     1  package api
     2  
     3  import (
     4  	"bytes"
     5  	"encoding/base64"
     6  	"fmt"
     7  	"net/http"
     8  
     9  	"github.com/diamondburned/arikawa/utils/json"
    10  	"github.com/pkg/errors"
    11  )
    12  
    13  var ErrInvalidImageCT = errors.New("unknown image content-type")
    14  var ErrInvalidImageData = errors.New("invalid image data")
    15  
    16  type ErrImageTooLarge struct {
    17  	Size, Max int
    18  }
    19  
    20  func (err ErrImageTooLarge) Error() string {
    21  	return fmt.Sprintf("Image is %.02fkb, larger than %.02fkb",
    22  		float64(err.Size)/1000, float64(err.Max)/1000)
    23  }
    24  
    25  // Image wraps around the Data URI Scheme that Discord uses:
    26  // https://discordapp.com/developers/docs/reference#image-data
    27  type Image struct {
    28  	// ContentType is optional and will be automatically detected. However, it
    29  	// should always return "image/jpeg," "image/png" or "image/gif".
    30  	ContentType string
    31  
    32  	// Just raw content of the file.
    33  	Content []byte
    34  }
    35  
    36  func DecodeImage(data []byte) (*Image, error) {
    37  	parts := bytes.SplitN(data, []byte{';'}, 2)
    38  	if len(parts) < 2 {
    39  		return nil, ErrInvalidImageData
    40  	}
    41  
    42  	if !bytes.HasPrefix(parts[0], []byte("data:")) {
    43  		return nil, errors.Wrap(ErrInvalidImageData, "invalid header")
    44  	}
    45  
    46  	if !bytes.HasPrefix(parts[1], []byte("base64,")) {
    47  		return nil, errors.Wrap(ErrInvalidImageData, "invalid base64")
    48  	}
    49  
    50  	var b64 = parts[1][len("base64,"):]
    51  	var img = Image{
    52  		ContentType: string(parts[0][len("data:"):]),
    53  		Content:     make([]byte, base64.StdEncoding.DecodedLen(len(b64))),
    54  	}
    55  
    56  	base64.StdEncoding.Decode(img.Content, b64)
    57  	return &img, nil
    58  }
    59  
    60  func (i Image) Validate(maxSize int) error {
    61  	if maxSize > 0 && len(i.Content) > maxSize {
    62  		return ErrImageTooLarge{len(i.Content), maxSize}
    63  	}
    64  
    65  	switch i.ContentType {
    66  	case "image/png", "image/jpeg", "image/gif":
    67  		return nil
    68  	default:
    69  		return ErrInvalidImageCT
    70  	}
    71  }
    72  
    73  func (i Image) Encode() ([]byte, error) {
    74  	if i.ContentType == "" {
    75  		var max = 512
    76  		if len(i.Content) < max {
    77  			max = len(i.Content)
    78  		}
    79  		i.ContentType = http.DetectContentType(i.Content[:max])
    80  	}
    81  
    82  	if err := i.Validate(0); err != nil {
    83  		return nil, err
    84  	}
    85  
    86  	b64enc := make([]byte, base64.StdEncoding.EncodedLen(len(i.Content)))
    87  	base64.StdEncoding.Encode(b64enc, i.Content)
    88  
    89  	return bytes.Join([][]byte{
    90  		[]byte("data:"),
    91  		[]byte(i.ContentType),
    92  		[]byte(";base64,"),
    93  		b64enc,
    94  	}, nil), nil
    95  }
    96  
    97  var _ json.Marshaler = (*Image)(nil)
    98  var _ json.Unmarshaler = (*Image)(nil)
    99  
   100  func (i Image) MarshalJSON() ([]byte, error) {
   101  	if len(i.Content) == 0 {
   102  		return []byte("null"), nil
   103  	}
   104  
   105  	b, err := i.Encode()
   106  	if err != nil {
   107  		return nil, err
   108  	}
   109  
   110  	return bytes.Join([][]byte{{'"'}, b, {'"'}}, nil), nil
   111  }
   112  
   113  func (i *Image) UnmarshalJSON(v []byte) error {
   114  	// Trim string
   115  	v = bytes.Trim(v, `"`)
   116  
   117  	// Accept a nil image.
   118  	if string(v) == "null" {
   119  		return nil
   120  	}
   121  
   122  	img, err := DecodeImage(v)
   123  	if err != nil {
   124  		return err
   125  	}
   126  
   127  	*i = *img
   128  	return nil
   129  }