jiajia
3 months ago
10 changed files with 789 additions and 39 deletions
@ -0,0 +1,193 @@ |
|||
package server |
|||
|
|||
import ( |
|||
"context" |
|||
"crypto/sha256" |
|||
"encoding/hex" |
|||
"encoding/json" |
|||
"fmt" |
|||
"io" |
|||
"net/url" |
|||
"os" |
|||
"path/filepath" |
|||
"strings" |
|||
"time" |
|||
) |
|||
|
|||
type mediaRecord struct { |
|||
media |
|||
UserID string `json:"user_id"` |
|||
LocalPath string `json:"local_path"` |
|||
} |
|||
|
|||
func mediaDir(cfg Config) string { |
|||
return filepath.Join(cfg.DataDir, "media") |
|||
} |
|||
|
|||
func mediaBlobDir(cfg Config) string { |
|||
return filepath.Join(mediaDir(cfg), "blobs") |
|||
} |
|||
|
|||
func mediaMetaDir(cfg Config) string { |
|||
return filepath.Join(mediaDir(cfg), "meta") |
|||
} |
|||
|
|||
func mediaMetaPath(cfg Config, mediaID string) string { |
|||
return filepath.Join(mediaMetaDir(cfg), mediaID+".json") |
|||
} |
|||
|
|||
func mediaBlobPath(cfg Config, mediaID, filename string) string { |
|||
ext := filepath.Ext(filename) |
|||
if ext == "" { |
|||
ext = extensionFromContentType("") |
|||
} |
|||
return filepath.Join(mediaBlobDir(cfg), mediaID+ext) |
|||
} |
|||
|
|||
func publicBaseURL(cfg Config) string { |
|||
base := strings.TrimSpace(cfg.PublicBaseURL) |
|||
if base == "" { |
|||
base = strings.TrimSpace(os.Getenv("POPIART_SERVER_ADDR")) |
|||
} |
|||
if base == "" { |
|||
base = "127.0.0.1:8080" |
|||
} |
|||
if !strings.HasPrefix(base, "http://") && !strings.HasPrefix(base, "https://") { |
|||
base = "http://" + base |
|||
} |
|||
base = strings.Replace(base, "0.0.0.0", "127.0.0.1", 1) |
|||
return strings.TrimRight(base, "/") |
|||
} |
|||
|
|||
func mediaContentURL(cfg Config, mediaID string) string { |
|||
return publicBaseURL(cfg) + "/v1/media/" + url.PathEscape(strings.TrimSpace(mediaID)) + "/content" |
|||
} |
|||
|
|||
func persistMediaContent(cfg Config, userID, projectID, artifactID, filename, contentType, visibility string, content []byte, mediaID string) (mediaRecord, error) { |
|||
if len(content) == 0 { |
|||
return mediaRecord{}, fmt.Errorf("media content is required") |
|||
} |
|||
filename = sanitizeFilename(strings.TrimSpace(filename)) |
|||
if filename == "" || filename == "artifact" { |
|||
filename = "artifact" + extensionFromContentType(contentType) |
|||
} |
|||
contentType = strings.TrimSpace(contentType) |
|||
if contentType == "" { |
|||
contentType = "application/octet-stream" |
|||
} |
|||
visibility = strings.TrimSpace(visibility) |
|||
if visibility == "" { |
|||
visibility = "unlisted" |
|||
} |
|||
mediaID = strings.TrimSpace(mediaID) |
|||
if mediaID == "" { |
|||
mediaID = "med_" + randomID() |
|||
} |
|||
if err := os.MkdirAll(mediaBlobDir(cfg), 0o755); err != nil { |
|||
return mediaRecord{}, err |
|||
} |
|||
if err := os.MkdirAll(mediaMetaDir(cfg), 0o755); err != nil { |
|||
return mediaRecord{}, err |
|||
} |
|||
|
|||
sum := sha256.Sum256(content) |
|||
blobPath := mediaBlobPath(cfg, mediaID, filename) |
|||
if err := os.WriteFile(blobPath, content, 0o644); err != nil { |
|||
return mediaRecord{}, err |
|||
} |
|||
|
|||
record := mediaRecord{ |
|||
media: media{ |
|||
ID: mediaID, |
|||
ArtifactID: strings.TrimSpace(artifactID), |
|||
ProjectID: strings.TrimSpace(projectID), |
|||
Filename: filename, |
|||
ContentType: contentType, |
|||
SizeBytes: int64(len(content)), |
|||
CreatedAt: time.Now().UTC().Format(time.RFC3339), |
|||
URL: mediaContentURL(cfg, mediaID), |
|||
Visibility: visibility, |
|||
SHA256: hex.EncodeToString(sum[:]), |
|||
}, |
|||
UserID: strings.TrimSpace(userID), |
|||
LocalPath: blobPath, |
|||
} |
|||
|
|||
metaBytes, err := json.Marshal(record) |
|||
if err != nil { |
|||
return mediaRecord{}, err |
|||
} |
|||
if err := os.WriteFile(mediaMetaPath(cfg, mediaID), metaBytes, 0o644); err != nil { |
|||
return mediaRecord{}, err |
|||
} |
|||
return record, nil |
|||
} |
|||
|
|||
func loadMediaRecord(cfg Config, mediaID string) (*mediaRecord, bool, error) { |
|||
mediaID = strings.TrimSpace(mediaID) |
|||
if mediaID == "" { |
|||
return nil, false, nil |
|||
} |
|||
data, err := os.ReadFile(mediaMetaPath(cfg, mediaID)) |
|||
if os.IsNotExist(err) { |
|||
return nil, false, nil |
|||
} |
|||
if err != nil { |
|||
return nil, false, err |
|||
} |
|||
var record mediaRecord |
|||
if err := json.Unmarshal(data, &record); err != nil { |
|||
return nil, false, err |
|||
} |
|||
return &record, true, nil |
|||
} |
|||
|
|||
func (s *Server) persistResultRefs(ctx context.Context, record *job, refs []resultRef) ([]resultRef, error) { |
|||
if record == nil || len(refs) == 0 { |
|||
return refs, nil |
|||
} |
|||
items := make([]resultRef, 0, len(refs)) |
|||
for idx, ref := range refs { |
|||
contentType, _, reader, err := s.newapi.openResultRef(ctx, record.UpstreamKey, ref) |
|||
if err != nil { |
|||
return nil, err |
|||
} |
|||
content, readErr := io.ReadAll(reader) |
|||
reader.Close() |
|||
if readErr != nil { |
|||
return nil, readErr |
|||
} |
|||
filename := strings.TrimSpace(ref.Filename) |
|||
if filename == "" { |
|||
filename = inferFilenameFromRef(ref, buildArtifactID(record.JobID, idx)) |
|||
} |
|||
artifactID := buildArtifactID(record.JobID, idx) |
|||
mediaRecord, err := persistMediaContent( |
|||
s.cfg, |
|||
record.UserID, |
|||
record.ProjectID, |
|||
artifactID, |
|||
filename, |
|||
defaultString(strings.TrimSpace(contentType), strings.TrimSpace(ref.ContentType)), |
|||
defaultString(strings.TrimSpace(ref.Visibility), "unlisted"), |
|||
content, |
|||
"", |
|||
) |
|||
if err != nil { |
|||
return nil, err |
|||
} |
|||
items = append(items, resultRef{ |
|||
Kind: "local_path", |
|||
URL: mediaRecord.URL, |
|||
LocalPath: mediaRecord.LocalPath, |
|||
MediaID: mediaRecord.ID, |
|||
Filename: mediaRecord.Filename, |
|||
ContentType: mediaRecord.ContentType, |
|||
SizeBytes: mediaRecord.SizeBytes, |
|||
Visibility: mediaRecord.Visibility, |
|||
SHA256: mediaRecord.SHA256, |
|||
StorageStatus: "ready", |
|||
}) |
|||
} |
|||
return items, nil |
|||
} |
|||
@ -0,0 +1,106 @@ |
|||
package server |
|||
|
|||
import ( |
|||
"bytes" |
|||
"encoding/json" |
|||
"io" |
|||
"mime/multipart" |
|||
"net/http" |
|||
"net/http/httptest" |
|||
"path/filepath" |
|||
"testing" |
|||
) |
|||
|
|||
func TestMediaUploadGetAndContent(t *testing.T) { |
|||
cfg := Config{ |
|||
SQLitePath: filepath.Join(t.TempDir(), "popiart.db"), |
|||
SkillhubDir: makeEmptySkillhub(t), |
|||
SessionSecret: "test-secret", |
|||
} |
|||
server, err := NewWithConfig(cfg) |
|||
if err != nil { |
|||
t.Fatalf("NewWithConfig: %v", err) |
|||
} |
|||
srv := httptest.NewServer(server.Handler()) |
|||
defer srv.Close() |
|||
server.cfg.PublicBaseURL = srv.URL |
|||
|
|||
sessionToken, _, ok, err := server.store.createSession("sk-media-user") |
|||
if err != nil { |
|||
t.Fatalf("createSession: %v", err) |
|||
} |
|||
if !ok { |
|||
t.Fatal("expected session creation to succeed") |
|||
} |
|||
|
|||
imageBytes := tinyPNG(t) |
|||
var body bytes.Buffer |
|||
writer := multipart.NewWriter(&body) |
|||
if err := writer.WriteField("project_id", "proj_media_demo"); err != nil { |
|||
t.Fatalf("write project_id field: %v", err) |
|||
} |
|||
if err := writer.WriteField("visibility", "public"); err != nil { |
|||
t.Fatalf("write visibility field: %v", err) |
|||
} |
|||
part, err := writer.CreateFormFile("file", "poster.png") |
|||
if err != nil { |
|||
t.Fatalf("create form file: %v", err) |
|||
} |
|||
if _, err := part.Write(imageBytes); err != nil { |
|||
t.Fatalf("write form file: %v", err) |
|||
} |
|||
if err := writer.Close(); err != nil { |
|||
t.Fatalf("close multipart writer: %v", err) |
|||
} |
|||
|
|||
req := httptest.NewRequest(http.MethodPost, "/v1/media/upload", &body) |
|||
req.Header.Set("Authorization", "Bearer "+sessionToken) |
|||
req.Header.Set("Content-Type", writer.FormDataContentType()) |
|||
rec := httptest.NewRecorder() |
|||
server.Handler().ServeHTTP(rec, req) |
|||
|
|||
if rec.Code != http.StatusCreated { |
|||
t.Fatalf("expected status 201, got %d body=%s", rec.Code, rec.Body.String()) |
|||
} |
|||
|
|||
var envelope struct { |
|||
OK bool `json:"ok"` |
|||
Data media `json:"data"` |
|||
} |
|||
if err := json.Unmarshal(rec.Body.Bytes(), &envelope); err != nil { |
|||
t.Fatalf("decode upload response: %v body=%s", err, rec.Body.String()) |
|||
} |
|||
if !envelope.OK { |
|||
t.Fatalf("expected ok upload response, got %s", rec.Body.String()) |
|||
} |
|||
if envelope.Data.ID == "" { |
|||
t.Fatalf("expected media id, got %#v", envelope.Data) |
|||
} |
|||
if envelope.Data.URL == "" { |
|||
t.Fatalf("expected stable url, got %#v", envelope.Data) |
|||
} |
|||
|
|||
getReq := httptest.NewRequest(http.MethodGet, "/v1/media/"+envelope.Data.ID, nil) |
|||
getReq.Header.Set("Authorization", "Bearer "+sessionToken) |
|||
getRec := httptest.NewRecorder() |
|||
server.Handler().ServeHTTP(getRec, getReq) |
|||
if getRec.Code != http.StatusOK { |
|||
t.Fatalf("expected media get 200, got %d body=%s", getRec.Code, getRec.Body.String()) |
|||
} |
|||
|
|||
contentResp, err := http.Get(envelope.Data.URL) |
|||
if err != nil { |
|||
t.Fatalf("GET media content: %v", err) |
|||
} |
|||
defer contentResp.Body.Close() |
|||
if contentResp.StatusCode != http.StatusOK { |
|||
t.Fatalf("expected media content 200, got %d", contentResp.StatusCode) |
|||
} |
|||
content, err := io.ReadAll(contentResp.Body) |
|||
if err != nil { |
|||
t.Fatalf("read content body: %v", err) |
|||
} |
|||
if !bytes.Equal(content, imageBytes) { |
|||
t.Fatal("expected media content bytes to match uploaded file") |
|||
} |
|||
} |
|||
Loading…
Reference in new issue