139 lines
3.1 KiB
Go
139 lines
3.1 KiB
Go
package ocrapi
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"io"
|
|
"mime/multipart"
|
|
"net/http"
|
|
"strings"
|
|
"time"
|
|
"unicode/utf8"
|
|
)
|
|
|
|
const (
|
|
defaultTimeout = 5 * time.Second
|
|
maximumReplyBytes = 64 << 10
|
|
)
|
|
|
|
var ErrServiceInvalid = errors.New("OCR service is invalid")
|
|
|
|
type Client struct {
|
|
endpoint string
|
|
http *http.Client
|
|
}
|
|
|
|
func NewClient(endpoint string, timeout time.Duration) (*Client, error) {
|
|
endpoint = strings.TrimSpace(endpoint)
|
|
if endpoint == "" {
|
|
return nil, nil
|
|
}
|
|
if timeout <= 0 {
|
|
timeout = defaultTimeout
|
|
}
|
|
return &Client{
|
|
endpoint: endpoint,
|
|
http: &http.Client{
|
|
Timeout: timeout,
|
|
CheckRedirect: func(*http.Request, []*http.Request) error {
|
|
return http.ErrUseLastResponse
|
|
},
|
|
},
|
|
}, nil
|
|
}
|
|
|
|
func (client *Client) Recognize(
|
|
ctx context.Context,
|
|
image []byte,
|
|
contentType string,
|
|
) (string, error) {
|
|
if client == nil || len(image) == 0 || !strings.HasPrefix(contentType, "image/") {
|
|
return "", ErrServiceInvalid
|
|
}
|
|
var body bytes.Buffer
|
|
writer := multipart.NewWriter(&body)
|
|
part, err := writer.CreateFormFile("file", "captcha"+extension(contentType))
|
|
if err != nil {
|
|
return "", ErrServiceInvalid
|
|
}
|
|
if _, err := part.Write(image); err != nil || writer.Close() != nil {
|
|
return "", ErrServiceInvalid
|
|
}
|
|
request, err := http.NewRequestWithContext(
|
|
ctx,
|
|
http.MethodPost,
|
|
client.endpoint,
|
|
&body,
|
|
)
|
|
if err != nil {
|
|
return "", ErrServiceInvalid
|
|
}
|
|
request.Header.Set("Content-Type", writer.FormDataContentType())
|
|
response, err := client.http.Do(request)
|
|
if err != nil {
|
|
return "", ErrServiceInvalid
|
|
}
|
|
defer response.Body.Close()
|
|
if response.StatusCode < http.StatusOK || response.StatusCode >= http.StatusMultipleChoices {
|
|
return "", ErrServiceInvalid
|
|
}
|
|
reply, err := readBounded(response.Body, maximumReplyBytes)
|
|
if err != nil {
|
|
return "", ErrServiceInvalid
|
|
}
|
|
code, err := parseCode(reply, response.Header.Get("Content-Type"))
|
|
if err != nil {
|
|
return "", ErrServiceInvalid
|
|
}
|
|
return code, nil
|
|
}
|
|
|
|
func parseCode(reply []byte, contentType string) (string, error) {
|
|
value := ""
|
|
if strings.HasPrefix(strings.ToLower(contentType), "application/json") {
|
|
var object map[string]any
|
|
if err := json.Unmarshal(reply, &object); err != nil {
|
|
return "", ErrServiceInvalid
|
|
}
|
|
for _, key := range []string{"text", "result", "data"} {
|
|
if candidate, ok := object[key].(string); ok {
|
|
value = candidate
|
|
break
|
|
}
|
|
}
|
|
} else {
|
|
value = string(reply)
|
|
}
|
|
value = strings.TrimSpace(value)
|
|
if value == "" || len([]byte(value)) > 64 || !utf8.ValidString(value) {
|
|
return "", ErrServiceInvalid
|
|
}
|
|
for _, character := range value {
|
|
if character < 32 || character == 127 {
|
|
return "", ErrServiceInvalid
|
|
}
|
|
}
|
|
return value, nil
|
|
}
|
|
|
|
func extension(contentType string) string {
|
|
switch contentType {
|
|
case "image/png":
|
|
return ".png"
|
|
case "image/jpeg":
|
|
return ".jpg"
|
|
default:
|
|
return ".img"
|
|
}
|
|
}
|
|
|
|
func readBounded(reader io.Reader, maximum int64) ([]byte, error) {
|
|
result, err := io.ReadAll(io.LimitReader(reader, maximum+1))
|
|
if err != nil || int64(len(result)) > maximum {
|
|
return nil, ErrServiceInvalid
|
|
}
|
|
return result, nil
|
|
}
|