jiajia
4 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