Julien Chaumond
4945fa3cd2
* Fix incorrect doc comments in huggingface engine * add huggingface bucket scanning * Update client.go * use limit=5000 on bucket listing endpoints * regenerate man page * escape bucket file path per-segment instead of whole * escape file path in bucket finding link too * use separate client without overall timeout for bucket file downloads
391 lines
12 KiB
Go
391 lines
12 KiB
Go
package huggingface
|
||
|
||
import (
|
||
"encoding/json"
|
||
"errors"
|
||
"fmt"
|
||
"io"
|
||
"net/http"
|
||
"net/url"
|
||
"strings"
|
||
"time"
|
||
|
||
"github.com/trufflesecurity/trufflehog/v3/pkg/context"
|
||
)
|
||
|
||
// Maps for API and HTML paths
|
||
var apiPaths = map[string]string{
|
||
DATASET: DatasetsRoute,
|
||
MODEL: ModelsAPIRoute,
|
||
SPACE: SpacesRoute,
|
||
BUCKET: BucketsRoute,
|
||
}
|
||
|
||
var htmlPaths = map[string]string{
|
||
DATASET: DatasetsRoute,
|
||
MODEL: "",
|
||
SPACE: SpacesRoute,
|
||
BUCKET: BucketsRoute,
|
||
}
|
||
|
||
type Author struct {
|
||
Username string `json:"name"`
|
||
}
|
||
|
||
type Latest struct {
|
||
Raw string `json:"raw"`
|
||
}
|
||
|
||
type Data struct {
|
||
Latest Latest `json:"latest"`
|
||
}
|
||
|
||
type Event struct {
|
||
Type string `json:"type"`
|
||
Author Author `json:"author"`
|
||
CreatedAt string `json:"createdAt"`
|
||
Data Data `json:"data"`
|
||
ID string `json:"id"`
|
||
}
|
||
|
||
func (e Event) GetAuthor() string {
|
||
return e.Author.Username
|
||
}
|
||
|
||
func (e Event) GetCreatedAt() string {
|
||
return e.CreatedAt
|
||
}
|
||
|
||
func (e Event) GetID() string {
|
||
return e.ID
|
||
}
|
||
|
||
type RepoData struct {
|
||
FullName string `json:"name"`
|
||
ResourceType string `json:"type"`
|
||
}
|
||
|
||
type Discussion struct {
|
||
ID int `json:"num"`
|
||
IsPR bool `json:"isPullRequest"`
|
||
CreatedAt string `json:"createdAt"`
|
||
Title string `json:"title"`
|
||
Events []Event `json:"events"`
|
||
Repo RepoData `json:"repo"`
|
||
}
|
||
|
||
func (d Discussion) GetID() string {
|
||
return fmt.Sprint(d.ID)
|
||
}
|
||
|
||
func (d Discussion) GetTitle() string {
|
||
return d.Title
|
||
}
|
||
|
||
func (d Discussion) GetCreatedAt() string {
|
||
return d.CreatedAt
|
||
}
|
||
|
||
func (d Discussion) GetRepo() string {
|
||
return d.Repo.FullName
|
||
}
|
||
|
||
// GetDiscussionPath returns the path (ex: "/models/user/repo/discussions/1") for the discussion
|
||
func (d Discussion) GetDiscussionPath() string {
|
||
basePath := fmt.Sprintf("%s/%s/%s", d.GetRepo(), DiscussionsRoute, d.GetID())
|
||
if d.Repo.ResourceType == "model" {
|
||
return basePath
|
||
}
|
||
return fmt.Sprintf("%s/%s", getResourceHTMLPath(d.Repo.ResourceType), basePath)
|
||
}
|
||
|
||
// GetGitPath returns the path (ex: "/models/user/repo.git") for the repo's git directory
|
||
func (d Discussion) GetGitPath() string {
|
||
basePath := fmt.Sprintf("%s.git", d.GetRepo())
|
||
if d.Repo.ResourceType == "model" {
|
||
return basePath
|
||
}
|
||
return fmt.Sprintf("%s/%s", getResourceHTMLPath(d.Repo.ResourceType), basePath)
|
||
}
|
||
|
||
type DiscussionList struct {
|
||
Discussions []Discussion `json:"discussions"`
|
||
}
|
||
|
||
type Repo struct {
|
||
IsPrivate bool `json:"private"`
|
||
Owner string `json:"author"`
|
||
RepoID string `json:"id"`
|
||
}
|
||
|
||
type Bucket struct {
|
||
IsPrivate bool `json:"private"`
|
||
BucketID string `json:"id"`
|
||
}
|
||
|
||
type BucketFile struct {
|
||
Type string `json:"type"`
|
||
Path string `json:"path"`
|
||
Size int64 `json:"size"`
|
||
}
|
||
|
||
type HFClient struct {
|
||
BaseURL string
|
||
APIKey string
|
||
HTTPClient *http.Client
|
||
// downloadClient streams file bodies. Unlike HTTPClient it has no overall
|
||
// Timeout, because http.Client.Timeout also interrupts Response.Body reads —
|
||
// which would kill large/slow downloads handed to handlers.HandleFile.
|
||
downloadClient *http.Client
|
||
}
|
||
|
||
// NewHFClient creates a new HF client
|
||
func NewHFClient(baseURL, apiKey string, timeout time.Duration) *HFClient {
|
||
// Bound the connect + response-header phase (including redirect hops to the
|
||
// CDN) without capping the body read; overall cancellation comes from the
|
||
// request context. Built explicitly (rather than cloning DefaultTransport)
|
||
// so it stays a *http.Transport even when tests swap the default transport.
|
||
downloadTransport := &http.Transport{
|
||
Proxy: http.ProxyFromEnvironment,
|
||
ForceAttemptHTTP2: true,
|
||
MaxIdleConns: 100,
|
||
IdleConnTimeout: 90 * time.Second,
|
||
TLSHandshakeTimeout: 10 * time.Second,
|
||
ExpectContinueTimeout: 1 * time.Second,
|
||
ResponseHeaderTimeout: timeout,
|
||
}
|
||
|
||
return &HFClient{
|
||
BaseURL: baseURL,
|
||
APIKey: apiKey,
|
||
HTTPClient: &http.Client{
|
||
Timeout: timeout,
|
||
},
|
||
downloadClient: &http.Client{
|
||
Transport: downloadTransport,
|
||
},
|
||
}
|
||
}
|
||
|
||
// get makes a GET request to the Hugging Face API
|
||
// Note: not addressing rate limit, since it seems very permissive. (ex: "If \
|
||
// your account suddenly sends 10k requests then you’re likely to receive 503")
|
||
func (c *HFClient) get(ctx context.Context, url string, target interface{}) error {
|
||
ctx, cancel := context.WithTimeout(ctx, 10*time.Second)
|
||
defer cancel()
|
||
|
||
req, err := http.NewRequestWithContext(ctx, "GET", url, nil)
|
||
if err != nil {
|
||
return fmt.Errorf("failed to create HuggingFace API request: %w", err)
|
||
}
|
||
|
||
req.Header.Set("Authorization", "Bearer "+c.APIKey)
|
||
|
||
resp, err := c.HTTPClient.Do(req)
|
||
if err != nil {
|
||
return fmt.Errorf("failed to make request to HuggingFace API: %w", err)
|
||
}
|
||
|
||
if resp.StatusCode == http.StatusUnauthorized {
|
||
return errors.New("invalid API key")
|
||
}
|
||
|
||
if resp.StatusCode == http.StatusForbidden {
|
||
return errors.New("access to this repo is restricted and you are not in the authorized list. Visit the repository to ask for access")
|
||
}
|
||
|
||
defer func() { _ = resp.Body.Close() }()
|
||
|
||
return json.NewDecoder(resp.Body).Decode(target)
|
||
}
|
||
|
||
// GetRepo retrieves repo from the Hugging Face API
|
||
func (c *HFClient) GetRepo(ctx context.Context, repoName string, resourceType string) (Repo, error) {
|
||
var repo Repo
|
||
url, err := buildAPIURL(c.BaseURL, resourceType, repoName)
|
||
if err != nil {
|
||
return repo, err
|
||
}
|
||
err = c.get(ctx, url, &repo)
|
||
return repo, err
|
||
}
|
||
|
||
// ListDiscussions retrieves discussions from the Hugging Face API
|
||
func (c *HFClient) ListDiscussions(ctx context.Context, repoInfo repoInfo) (DiscussionList, error) {
|
||
var discussions DiscussionList
|
||
baseURL, err := buildAPIURL(c.BaseURL, string(repoInfo.resourceType), repoInfo.fullName)
|
||
if err != nil {
|
||
return discussions, err
|
||
}
|
||
url := fmt.Sprintf("%s/%s", baseURL, DiscussionsRoute)
|
||
err = c.get(ctx, url, &discussions)
|
||
return discussions, err
|
||
}
|
||
|
||
func (c *HFClient) GetDiscussionByID(ctx context.Context, repoInfo repoInfo, discussionID string) (Discussion, error) {
|
||
var discussion Discussion
|
||
baseURL, err := buildAPIURL(c.BaseURL, string(repoInfo.resourceType), repoInfo.fullName)
|
||
if err != nil {
|
||
return discussion, err
|
||
}
|
||
url := fmt.Sprintf("%s/%s/%s", baseURL, DiscussionsRoute, discussionID)
|
||
err = c.get(ctx, url, &discussion)
|
||
return discussion, err
|
||
}
|
||
|
||
// ListReposByAuthor retrieves repos from the Hugging Face API by author (user or org)
|
||
// Note: not addressing pagination b/c allow by default 1000 results, which should be enough for 99.99% of cases
|
||
func (c *HFClient) ListReposByAuthor(ctx context.Context, resourceType string, author string) ([]Repo, error) {
|
||
var repos []Repo
|
||
url := fmt.Sprintf("%s/%s/%s?limit=1000&author=%s", c.BaseURL, APIRoute, getResourceAPIPath(resourceType), author)
|
||
err := c.get(ctx, url, &repos)
|
||
return repos, err
|
||
}
|
||
|
||
// GetBucket retrieves bucket metadata from the Hugging Face API.
|
||
func (c *HFClient) GetBucket(ctx context.Context, bucketID string) (Bucket, error) {
|
||
var bucket Bucket
|
||
url, err := buildAPIURL(c.BaseURL, BUCKET, bucketID)
|
||
if err != nil {
|
||
return bucket, err
|
||
}
|
||
err = c.get(ctx, url, &bucket)
|
||
return bucket, err
|
||
}
|
||
|
||
// ListBucketsByAuthor retrieves buckets from the Hugging Face API by author (user or org).
|
||
func (c *HFClient) ListBucketsByAuthor(ctx context.Context, author string) ([]Bucket, error) {
|
||
var buckets []Bucket
|
||
url := fmt.Sprintf("%s/%s/%s/%s?limit=%d", c.BaseURL, APIRoute, BucketsRoute, author, bucketsPagination)
|
||
for url != "" {
|
||
var page []Bucket
|
||
next, err := c.getPage(ctx, url, &page)
|
||
if err != nil {
|
||
return buckets, err
|
||
}
|
||
buckets = append(buckets, page...)
|
||
url = next
|
||
}
|
||
return buckets, nil
|
||
}
|
||
|
||
// ListBucketFiles recursively lists all files in a bucket, following Link
|
||
// header pagination.
|
||
func (c *HFClient) ListBucketFiles(ctx context.Context, bucketID string) ([]BucketFile, error) {
|
||
var files []BucketFile
|
||
url := fmt.Sprintf("%s/%s/%s/%s/tree?recursive=true&limit=%d", c.BaseURL, APIRoute, BucketsRoute, bucketID, bucketsPagination)
|
||
for url != "" {
|
||
var page []BucketFile
|
||
next, err := c.getPage(ctx, url, &page)
|
||
if err != nil {
|
||
return files, err
|
||
}
|
||
files = append(files, page...)
|
||
url = next
|
||
}
|
||
return files, nil
|
||
}
|
||
|
||
// DownloadBucketFile streams a file's content from a bucket via the resolve
|
||
// endpoint. The caller is responsible for closing the returned reader.
|
||
func (c *HFClient) DownloadBucketFile(ctx context.Context, bucketID string, path string) (io.ReadCloser, error) {
|
||
// Escape each path segment individually so special characters (spaces,
|
||
// '?', '#', etc.) are encoded while slashes are preserved as separators.
|
||
downloadURL := fmt.Sprintf("%s/%s/%s/resolve/%s", c.BaseURL, BucketsRoute, bucketID, escapePathSegments(path))
|
||
|
||
req, err := http.NewRequestWithContext(ctx, "GET", downloadURL, nil)
|
||
if err != nil {
|
||
return nil, fmt.Errorf("failed to create HuggingFace bucket download request: %w", err)
|
||
}
|
||
req.Header.Set("Authorization", "Bearer "+c.APIKey)
|
||
|
||
resp, err := c.downloadClient.Do(req)
|
||
if err != nil {
|
||
return nil, fmt.Errorf("failed to download bucket file %s: %w", path, err)
|
||
}
|
||
if resp.StatusCode != http.StatusOK {
|
||
_ = resp.Body.Close()
|
||
return nil, fmt.Errorf("failed to download bucket file %s: status %d", path, resp.StatusCode)
|
||
}
|
||
return resp.Body, nil
|
||
}
|
||
|
||
// getPage makes a GET request like get, but additionally returns the "next"
|
||
// URL from the response Link header, if any, for paginated endpoints.
|
||
func (c *HFClient) getPage(ctx context.Context, url string, target interface{}) (string, error) {
|
||
ctx, cancel := context.WithTimeout(ctx, 10*time.Second)
|
||
defer cancel()
|
||
|
||
req, err := http.NewRequestWithContext(ctx, "GET", url, nil)
|
||
if err != nil {
|
||
return "", fmt.Errorf("failed to create HuggingFace API request: %w", err)
|
||
}
|
||
req.Header.Set("Authorization", "Bearer "+c.APIKey)
|
||
|
||
resp, err := c.HTTPClient.Do(req)
|
||
if err != nil {
|
||
return "", fmt.Errorf("failed to make request to HuggingFace API: %w", err)
|
||
}
|
||
defer func() { _ = resp.Body.Close() }()
|
||
|
||
if resp.StatusCode == http.StatusUnauthorized {
|
||
return "", errors.New("invalid API key")
|
||
}
|
||
if resp.StatusCode == http.StatusForbidden {
|
||
return "", errors.New("access to this resource is restricted and you are not in the authorized list. Visit the resource to ask for access")
|
||
}
|
||
if resp.StatusCode != http.StatusOK {
|
||
return "", fmt.Errorf("unexpected status %d from HuggingFace API", resp.StatusCode)
|
||
}
|
||
|
||
if err := json.NewDecoder(resp.Body).Decode(target); err != nil {
|
||
return "", err
|
||
}
|
||
return parseNextLink(resp.Header.Get("Link")), nil
|
||
}
|
||
|
||
// escapePathSegments URL-encodes each slash-separated segment of a path while
|
||
// keeping the slashes as separators, e.g. "a b/c?d" -> "a%20b/c%3Fd".
|
||
func escapePathSegments(path string) string {
|
||
segments := strings.Split(path, "/")
|
||
for i, segment := range segments {
|
||
segments[i] = url.PathEscape(segment)
|
||
}
|
||
return strings.Join(segments, "/")
|
||
}
|
||
|
||
// parseNextLink extracts the URL with rel="next" from a Link header value.
|
||
// Example: <https://huggingface.co/api/...?cursor=abc>; rel="next"
|
||
func parseNextLink(header string) string {
|
||
for _, part := range strings.Split(header, ",") {
|
||
sections := strings.Split(part, ";")
|
||
if len(sections) < 2 {
|
||
continue
|
||
}
|
||
urlPart := strings.Trim(strings.TrimSpace(sections[0]), "<>")
|
||
for _, param := range sections[1:] {
|
||
if strings.TrimSpace(param) == `rel="next"` {
|
||
return urlPart
|
||
}
|
||
}
|
||
}
|
||
return ""
|
||
}
|
||
|
||
// getResourceAPIPath returns the API path for the given resource type
|
||
func getResourceAPIPath(resourceType string) string {
|
||
return apiPaths[resourceType]
|
||
}
|
||
|
||
// getResourceHTMLPath returns the HTML path for the given resource type
|
||
func getResourceHTMLPath(resourceType string) string {
|
||
return htmlPaths[resourceType]
|
||
}
|
||
|
||
func buildAPIURL(endpoint string, resourceType string, repoName string) (string, error) {
|
||
if endpoint == "" || resourceType == "" || repoName == "" {
|
||
return "", errors.New("endpoint, resourceType, and repoName must not be empty")
|
||
}
|
||
return fmt.Sprintf("%s/%s/%s/%s", endpoint, APIRoute, getResourceAPIPath(resourceType), repoName), nil
|
||
}
|