From c36ad9ced74ab5932fffb9a87dacfb2cbf92308e Mon Sep 17 00:00:00 2001 From: Sleroq Date: Sat, 15 Aug 2026 16:44:35 +0300 Subject: [PATCH 01/33] feat: add Pixiv media and bookmark support --- apps/server/package.json | 1 + apps/server/src/handlers/image.ts | 80 +++--- apps/server/src/handlers/pixiv.ts | 108 ++++++++ apps/server/src/index.ts | 27 +- apps/server/src/queue/absurd.ts | 2 + apps/server/src/queue/classification.ts | 14 +- apps/server/src/queue/embeddings.ts | 8 +- apps/server/src/queue/image-collector.ts | 236 ++++++++---------- apps/server/src/queue/pixiv.ts | 166 ++++++++++++ apps/server/src/queue/scrapper.ts | 68 ++--- apps/server/src/scripts/classification.ts | 14 +- apps/server/src/scripts/embeddings.ts | 12 +- .../src/scripts/update-photo-dimensions.ts | 13 +- .../src/services/duplicate-detection.ts | 5 +- apps/server/src/services/media-download.ts | 42 ++++ apps/server/src/services/pixiv-media.test.ts | 103 ++++++++ apps/server/src/services/pixiv-media.ts | 159 ++++++++++++ apps/web/src/routes/settings.tsx | 95 ++++++- bun.lock | 72 +++++- packages/api/package.json | 1 + packages/api/src/routers/index.ts | 6 + packages/api/src/routers/pixiv.ts | 57 +++++ packages/api/src/routers/profiles.ts | 9 + packages/api/src/routers/search.ts | 95 ++++--- packages/api/src/routers/tweets.ts | 30 ++- .../services/pixiv-credential-operation.ts | 46 ++++ packages/api/src/services/pixiv-credential.ts | 56 +++++ packages/api/src/services/pixiv.ts | 109 ++++++++ packages/api/src/types/tweets.ts | 9 + .../api/src/utils/search-transformations.ts | 73 ++++++ packages/api/src/utils/transformations.ts | 120 +++------ packages/api/tests/pixiv-credential.test.ts | 68 +++++ packages/api/tests/transformations.test.ts | 39 +++ packages/crypto/src/cookie-encryption.ts | 20 +- .../migration.sql | 65 +++++ packages/utils/prisma/schema.prisma | 88 ++++--- packages/utils/src/db.ts | 21 +- 37 files changed, 1719 insertions(+), 418 deletions(-) create mode 100644 apps/server/src/handlers/pixiv.ts create mode 100644 apps/server/src/queue/pixiv.ts create mode 100644 apps/server/src/services/media-download.ts create mode 100644 apps/server/src/services/pixiv-media.test.ts create mode 100644 apps/server/src/services/pixiv-media.ts create mode 100644 packages/api/src/routers/pixiv.ts create mode 100644 packages/api/src/services/pixiv-credential-operation.ts create mode 100644 packages/api/src/services/pixiv-credential.ts create mode 100644 packages/api/src/services/pixiv.ts create mode 100644 packages/api/src/utils/search-transformations.ts create mode 100644 packages/api/tests/pixiv-credential.test.ts create mode 100644 packages/api/tests/transformations.test.ts create mode 100644 packages/utils/prisma/migrations/20260814000000_provider_neutral_posts_pixiv/migration.sql diff --git a/apps/server/package.json b/apps/server/package.json index d928dbd7..2a8fb6bf 100644 --- a/apps/server/package.json +++ b/apps/server/package.json @@ -39,6 +39,7 @@ "date-fns": "^4.4.0", "effect": "4.0.0-beta.100", "grammy": "catalog:", + "jszip": "^3.10.1", "pg": "^8.22.0", "pino": "^10.3.1", "pino-pretty": "^13.1.3", diff --git a/apps/server/src/handlers/image.ts b/apps/server/src/handlers/image.ts index 425bd592..46628be7 100644 --- a/apps/server/src/handlers/image.ts +++ b/apps/server/src/handlers/image.ts @@ -19,6 +19,7 @@ type InlineImageSearchResult = { photo_id: string; s3_path: string; tweet_id: string; + source_url: string; username: string | null; height: number | null; width: number | null; @@ -72,7 +73,7 @@ async function searchInlineImagesWithLegacyQuery( while (allPhotos.length < photoOffset + pageQueryLimit) { const { authors, textQuery } = parseInlineImageQuery(query); - const whereClause: Prisma.TweetWhereInput = {}; + const whereClause: Prisma.PostWhereInput = {}; if (authors.length > 0 && textQuery) { whereClause.AND = [ @@ -95,7 +96,7 @@ async function searchInlineImagesWithLegacyQuery( logger, { searchMode: "legacy", userId, tweetSkip, pageSize: INLINE_QUERY_PAGE_SIZE }, () => - prisma.tweet.findMany({ + prisma.post.findMany({ where: { userId, photos: { @@ -142,6 +143,7 @@ async function searchInlineImagesWithLegacyQuery( photo_id: photo.id, s3_path: photo.s3Path as string, tweet_id: tweet.id, + source_url: tweet.sourceUrl, username: tweet.username, height: photo.height, width: photo.width, @@ -225,7 +227,7 @@ composer.on("inline_query").filter( 200, ); const queryTime = new Date().toISOString(); - const photoDedupeKey = Prisma.sql`COALESCE(NULLIF(p.perceptual_hash, ''), p.id)`; + const photoDedupeKey = Prisma.sql`COALESCE(NULLIF(p.perceptual_hash, ''), p.provider || ':' || p.external_id)`; const authorFilter = authors.length > 0 @@ -271,10 +273,10 @@ composer.on("inline_query").filter( OR lower(general_tag.value) LIKE ${queryStartsWith} OR lower(general_tag.value) LIKE ${queryContains} ) - OR lower(COALESCE(t.tweet_text, '')) LIKE ${queryContains} + OR lower(COALESCE(t.text, '')) LIKE ${queryContains} OR EXISTS ( SELECT 1 - FROM jsonb_array_elements_text(COALESCE(t.tweet_data->'hashtags', '[]'::jsonb)) AS hashtag(value) + FROM jsonb_array_elements_text(COALESCE(t.provider_payload->'hashtags', '[]'::jsonb)) AS hashtag(value) WHERE lower(hashtag.value) = ${queryLower} OR lower(hashtag.value) LIKE ${queryStartsWith} OR lower(hashtag.value) LIKE ${queryContains} @@ -334,7 +336,7 @@ composer.on("inline_query").filter( ELSE 0.0 END ) - FROM jsonb_array_elements_text(COALESCE(t.tweet_data->'hashtags', '[]'::jsonb)) AS hashtag(value) + FROM jsonb_array_elements_text(COALESCE(t.provider_payload->'hashtags', '[]'::jsonb)) AS hashtag(value) ), 0.0 ) @@ -342,7 +344,7 @@ composer.on("inline_query").filter( : Prisma.sql`0.0`; const tweetTextScore = hasTextQuery - ? Prisma.sql`CASE WHEN lower(COALESCE(t.tweet_text, '')) LIKE ${queryContains} THEN 0.34 ELSE 0.0 END` + ? Prisma.sql`CASE WHEN lower(COALESCE(t.text, '')) LIKE ${queryContains} THEN 0.34 ELSE 0.0 END` : Prisma.sql`0.0`; let rankedPhotos: InlineImageSearchResult[] = []; @@ -366,19 +368,20 @@ composer.on("inline_query").filter( prisma.$queryRaw(Prisma.sql` WITH ranked AS ( SELECT - p.id AS photo_id, + p.external_id AS photo_id, p.s3_path, - t.id AS tweet_id, + t.external_id AS tweet_id, + t.source_url, t.username, p.height, p.width, p.created_at AS photo_created_at, ROW_NUMBER() OVER ( PARTITION BY ${photoDedupeKey} - ORDER BY p.created_at DESC, p.id DESC + ORDER BY p.created_at DESC, p.external_id DESC ) AS duplicate_rank - FROM photos p - JOIN tweets t ON t.id = p.tweet_id AND t.user_id = p.user_id + FROM media p + JOIN posts t ON t.external_id = p.post_external_id AND t.user_id = p.user_id AND t.provider = p.provider WHERE p.user_id = ${userId} AND p.deleted_at IS NULL AND p.s3_path IS NOT NULL @@ -388,6 +391,7 @@ composer.on("inline_query").filter( photo_id, s3_path, tweet_id, + source_url, username, height, width, @@ -435,9 +439,9 @@ composer.on("inline_query").filter( () => prisma.$queryRaw(Prisma.sql` WITH image_candidates AS ( - SELECT p.id, p.user_id - FROM photos p - JOIN tweets t ON t.id = p.tweet_id AND t.user_id = p.user_id + SELECT p.external_id AS id, p.user_id, p.provider + FROM media p + JOIN posts t ON t.external_id = p.post_external_id AND t.user_id = p.user_id AND t.provider = p.provider WHERE p.user_id = ${userId} AND p.deleted_at IS NULL AND p.s3_path IS NOT NULL @@ -449,9 +453,9 @@ composer.on("inline_query").filter( LIMIT ${candidateLimit} ), tag_candidates AS ( - SELECT p.id, p.user_id - FROM photos p - JOIN tweets t ON t.id = p.tweet_id AND t.user_id = p.user_id + SELECT p.external_id AS id, p.user_id, p.provider + FROM media p + JOIN posts t ON t.external_id = p.post_external_id AND t.user_id = p.user_id AND t.provider = p.provider WHERE p.user_id = ${userId} AND p.deleted_at IS NULL AND p.s3_path IS NOT NULL @@ -463,9 +467,9 @@ composer.on("inline_query").filter( LIMIT ${candidateLimit} ), lexical_candidates AS ( - SELECT p.id, p.user_id - FROM photos p - JOIN tweets t ON t.id = p.tweet_id AND t.user_id = p.user_id + SELECT p.external_id AS id, p.user_id, p.provider + FROM media p + JOIN posts t ON t.external_id = p.post_external_id AND t.user_id = p.user_id AND t.provider = p.provider WHERE p.user_id = ${userId} AND p.deleted_at IS NULL AND p.s3_path IS NOT NULL @@ -474,24 +478,25 @@ composer.on("inline_query").filter( LIMIT ${candidateLimit} ), candidate_pool AS ( - SELECT DISTINCT id, user_id + SELECT DISTINCT id, user_id, provider FROM ( - SELECT id, user_id FROM image_candidates + SELECT id, user_id, provider FROM image_candidates UNION ALL - SELECT id, user_id FROM tag_candidates + SELECT id, user_id, provider FROM tag_candidates UNION ALL - SELECT id, user_id FROM lexical_candidates + SELECT id, user_id, provider FROM lexical_candidates ) candidates ), scored AS ( SELECT - p.id AS photo_id, + p.external_id AS photo_id, ${photoDedupeKey} AS dedupe_key, p.s3_path, p.height, p.width, t.username, - t.id AS tweet_id, + t.external_id AS tweet_id, + t.source_url, t.created_at AS tweet_created_at, COALESCE(1.0 - (p.image_vec <=> ${textVector}::vector), 0.0) AS s_image, COALESCE(1.0 - (p.tag_vec <=> ${textVector}::vector), 0.0) AS s_tag_semantic, @@ -501,8 +506,8 @@ composer.on("inline_query").filter( ${tweetTextScore} AS s_tweet_text, ${authorScore} AS s_author FROM candidate_pool c - JOIN photos p ON p.id = c.id AND p.user_id = c.user_id - JOIN tweets t ON t.id = p.tweet_id AND t.user_id = p.user_id + JOIN media p ON p.external_id = c.id AND p.user_id = c.user_id AND p.provider = c.provider + JOIN posts t ON t.external_id = p.post_external_id AND t.user_id = p.user_id AND t.provider = p.provider ), fused AS ( SELECT @@ -539,7 +544,7 @@ composer.on("inline_query").filter( ) AS duplicate_rank FROM fused ) - SELECT photo_id, s3_path, tweet_id, username, height, width, final_score + SELECT photo_id, s3_path, tweet_id, source_url, username, height, width, final_score FROM deduped WHERE duplicate_rank = 1 ORDER BY final_score DESC NULLS LAST, photo_id DESC @@ -557,21 +562,22 @@ composer.on("inline_query").filter( prisma.$queryRaw(Prisma.sql` WITH scored AS ( SELECT - p.id AS photo_id, + p.external_id AS photo_id, ${photoDedupeKey} AS dedupe_key, p.s3_path, p.height, p.width, t.username, - t.id AS tweet_id, + t.external_id AS tweet_id, + t.source_url, t.created_at AS tweet_created_at, ${characterScore} AS s_character, ${tagLexicalScore} AS s_tag_lexical, ${hashtagScore} AS s_hashtag, ${tweetTextScore} AS s_tweet_text, ${authorScore} AS s_author - FROM photos p - JOIN tweets t ON t.id = p.tweet_id AND t.user_id = p.user_id + FROM media p + JOIN posts t ON t.external_id = p.post_external_id AND t.user_id = p.user_id AND t.provider = p.provider WHERE p.user_id = ${userId} AND p.deleted_at IS NULL AND p.s3_path IS NOT NULL @@ -612,7 +618,7 @@ composer.on("inline_query").filter( ) AS duplicate_rank FROM fused ) - SELECT photo_id, s3_path, tweet_id, username, height, width, final_score + SELECT photo_id, s3_path, tweet_id, source_url, username, height, width, final_score FROM deduped WHERE duplicate_rank = 1 ORDER BY final_score DESC NULLS LAST, photo_id DESC @@ -648,8 +654,8 @@ composer.on("inline_query").filter( const results = photosForThisPage.map((photo) => { const photoUrl = `${env.BASE_CDN_URL}/${photo.s3_path}`; const caption = photo.username - ? FormattedString.link(`@${photo.username}`, `https://x.com/i/status/${photo.tweet_id}`) - : new FormattedString(`https://x.com/i/status/${photo.tweet_id}`); + ? FormattedString.link(`@${photo.username}`, photo.source_url) + : new FormattedString(photo.source_url); return InlineQueryResultBuilder.photo(photo.photo_id, photoUrl, { caption: caption.caption, diff --git a/apps/server/src/handlers/pixiv.ts b/apps/server/src/handlers/pixiv.ts new file mode 100644 index 00000000..d86190c8 --- /dev/null +++ b/apps/server/src/handlers/pixiv.ts @@ -0,0 +1,108 @@ +import { mkdtemp, rm, writeFile } from "node:fs/promises"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { withPixivClient } from "@starlight/api/services/pixiv-credential"; +import { Composer, InputFile } from "grammy"; +import { readResponseBounded } from "@/services/media-download"; +import { getScheduledPixivGeneration, pixivApp } from "@/queue/pixiv"; +import { RETRY } from "@/queue/absurd"; +import { + convertUgoira, + extractUgoiraZip, + MAX_PIXIV_DOWNLOAD_BYTES, + parsePixivArtworkUrl, +} from "@/services/pixiv-media"; +import type { Context } from "@/types"; + +const MAX_MANGA_BYTES = 150_000_000; +const pixivHandler = new Composer(); + +pixivHandler.command("pixiv", async (ctx) => { + if (!ctx.user) { + return; + } + const generation = getScheduledPixivGeneration(); + await pixivApp.spawn( + "scheduled-pixiv-bookmarks", + { generation, userId: ctx.user.id, limit: 300 }, + { + idempotencyKey: `scheduled-pixiv-${ctx.user.id}-${generation}`, + maxAttempts: 3, + retryStrategy: RETRY.pixiv, + }, + ); + await pixivApp.spawn( + "pixiv-bookmarks", + { userId: ctx.user.id, count: 0, limit: 300 }, + { maxAttempts: 3, retryStrategy: RETRY.pixiv }, + ); + await ctx.reply("Starting Pixiv bookmark sync. Check your gallery in a few minutes."); +}); + +pixivHandler.on("message:text", async (ctx, next) => { + const id = parsePixivArtworkUrl(ctx.message.text.trim()); + if (!id) { + return next(); + } + if (!ctx.user) { + return; + } + const directory = await mkdtemp(join(tmpdir(), "starlight-pixiv-")); + try { + const handled = await withPixivClient(ctx.user.id, async (client) => { + const artwork = await client.artwork(id); + if (artwork.type === "ugoira") { + const metadata = await client.ugoira(id); + const archive = await readResponseBounded(await client.fetchMedia(metadata.zipUrls.medium)); + const extracted = await extractUgoiraZip(archive, metadata.frames); + try { + const output = join(extracted.directory, "ugoira.mp4"); + await convertUgoira(extracted.concatPath, output); + await ctx.replyWithVideo(new InputFile(output), { + caption: artwork.title, + }); + } finally { + await rm(extracted.directory, { recursive: true, force: true }); + } + return true; + } + + let aggregateBytes = 0; + const files: InputFile[] = []; + for (const [position, url] of artwork.mediaUrls.entries()) { + const bytes = await readResponseBounded( + await client.fetchMedia(url), + Math.min(MAX_PIXIV_DOWNLOAD_BYTES, MAX_MANGA_BYTES - aggregateBytes), + ); + aggregateBytes += bytes.byteLength; + const extension = new URL(url).pathname.split(".").at(-1) ?? "jpg"; + const path = join(directory, `${position}.${extension}`); + await writeFile(path, bytes); + files.push(new InputFile(path)); + } + for (let offset = 0; offset < files.length; offset += 10) { + const chunk = files.slice(offset, offset + 10); + const single = chunk.at(0); + if (chunk.length === 1 && single) { + await ctx.replyWithDocument(single, { caption: artwork.title }); + } else { + await ctx.replyWithMediaGroup( + chunk.map((file, index) => ({ + type: "document" as const, + media: file, + caption: index === 0 ? artwork.title : undefined, + })), + ); + } + } + return true; + }); + if (handled === undefined) { + await ctx.reply("Connect Pixiv in Settings first."); + } + } finally { + await rm(directory, { recursive: true, force: true }); + } +}); + +export default pixivHandler; diff --git a/apps/server/src/index.ts b/apps/server/src/index.ts index c87ba430..2f152456 100644 --- a/apps/server/src/index.ts +++ b/apps/server/src/index.ts @@ -5,6 +5,7 @@ import "@/services/runtime"; import chatMemberHandler from "@/handlers/chat-member"; import imageHandler from "@/handlers/image"; import messageHandler from "@/handlers/message"; +import pixivHandler from "@/handlers/pixiv"; import startHandler from "@/handlers/start"; import tweetImageHandler from "@/handlers/tweet-image"; import videoHandler from "@/handlers/video"; @@ -14,6 +15,7 @@ import { classificationApp } from "@/queue/classification"; import { embeddingsApp } from "@/queue/embeddings"; import { imagesApp } from "@/queue/image-collector"; import { memoryApp } from "@/queue/memory"; +import { pixivApp } from "@/queue/pixiv"; import { scrapperApp } from "@/queue/scrapper"; initTelemetry(); @@ -31,6 +33,7 @@ const boundary = bot.errorBoundary((error) => { }); boundary.use(videoHandler); +boundary.use(pixivHandler); boundary.use(tweetImageHandler); boundary.use(imageHandler); boundary.use(messageHandler); @@ -38,13 +41,15 @@ boundary.use(startHandler); boundary.use(chatMemberHandler); await Promise.all( - [imagesApp, classificationApp, embeddingsApp, scrapperApp, memoryApp].map(async (app) => { - await app.createQueue(); - await app.setQueuePolicy(undefined, { - cleanupLimit: 2000, - cleanupTtl: "1 day", - }); - }), + [imagesApp, classificationApp, embeddingsApp, scrapperApp, memoryApp, pixivApp].map( + async (app) => { + await app.createQueue(); + await app.setQueuePolicy(undefined, { + cleanupLimit: 2000, + cleanupTtl: "1 day", + }); + }, + ), ); const workers = await Promise.all([ imagesApp.startWorker({ @@ -80,8 +85,14 @@ const workers = await Promise.all([ onError: (error) => logger.error({ err: error }, "Memory worker error"), workerId: QUEUES.memory, }), + pixivApp.startWorker({ + batchSize: 1, + concurrency: 1, + onError: (error) => logger.error({ err: error }, "Pixiv worker error"), + workerId: QUEUES.pixiv, + }), ]); -const queueApps = [imagesApp, classificationApp, embeddingsApp, scrapperApp, memoryApp]; +const queueApps = [imagesApp, classificationApp, embeddingsApp, scrapperApp, memoryApp, pixivApp]; const runner = run(bot); logger.info("Bot is running..."); diff --git a/apps/server/src/queue/absurd.ts b/apps/server/src/queue/absurd.ts index 9a0e9842..f068cffe 100644 --- a/apps/server/src/queue/absurd.ts +++ b/apps/server/src/queue/absurd.ts @@ -42,6 +42,7 @@ export const QUEUES = { embeddings: "embeddings", images: "images-collector", memory: "chat-memory", + pixiv: "pixiv-bookmarks", scrapper: "feed-scrapper", } as const; @@ -50,5 +51,6 @@ export const RETRY = { embeddings: { kind: "exponential", baseSeconds: 30, factor: 2 } satisfies RetryStrategy, images: { kind: "exponential", baseSeconds: 10, factor: 2 } satisfies RetryStrategy, memory: { kind: "exponential", baseSeconds: 20, factor: 2 } satisfies RetryStrategy, + pixiv: { kind: "exponential", baseSeconds: 150, factor: 2 } satisfies RetryStrategy, scrapper: { kind: "exponential", baseSeconds: 150, factor: 2 } satisfies RetryStrategy, } as const; diff --git a/apps/server/src/queue/classification.ts b/apps/server/src/queue/classification.ts index 0c353348..00e5d4a0 100644 --- a/apps/server/src/queue/classification.ts +++ b/apps/server/src/queue/classification.ts @@ -8,6 +8,7 @@ import type { Classification } from "@/types"; interface ClassificationJobData { photoId: string; + provider?: string; requestId?: string; userId: string; } @@ -27,6 +28,7 @@ classificationApp.registerTask( } const { photoId, userId, requestId: incomingRequestId } = params; + const provider = params.provider ?? "twitter"; const requestId = incomingRequestId || Bun.randomUUIDv7(); if (!(env.ML_BASE_URL && env.ML_API_TOKEN)) { @@ -37,8 +39,8 @@ classificationApp.registerTask( logger.info({ photoId, userId, requestId }, "Classifying photo"); // Fetch photo record to get URL - const photo = await prisma.photo.findUnique({ - where: { photoId: { id: photoId, userId } }, + const photo = await prisma.media.findUnique({ + where: { mediaId: { id: photoId, userId, provider } }, select: { id: true, userId: true, @@ -101,16 +103,16 @@ classificationApp.registerTask( throw error; } - await prisma.photo.update({ - where: { photoId: { id: photoId, userId } }, + await prisma.media.update({ + where: { mediaId: { id: photoId, userId, provider } }, data: { classification: data }, }); await embeddingsApp.spawn( "embeddings", - { photoId, userId, requestId }, + { photoId, provider, userId, requestId }, { - idempotencyKey: `embed-${photoId}-${userId}`, + idempotencyKey: `embed-${provider}-${photoId}-${userId}`, maxAttempts: 5, retryStrategy: RETRY.embeddings, }, diff --git a/apps/server/src/queue/embeddings.ts b/apps/server/src/queue/embeddings.ts index 538c52dc..fdd5ed09 100644 --- a/apps/server/src/queue/embeddings.ts +++ b/apps/server/src/queue/embeddings.ts @@ -7,6 +7,7 @@ import { runtime } from "@/services/runtime"; interface ClassificationJobData { photoId: string; + provider?: string; requestId?: string; userId: string; } @@ -24,6 +25,7 @@ embeddingsApp.registerTask({ name: "embeddings" }, async } const { photoId, userId, requestId: incomingRequestId } = params; + const provider = params.provider ?? "twitter"; const requestId = incomingRequestId || Bun.randomUUIDv7(); if (!(env.ML_BASE_URL && env.ML_API_TOKEN)) { @@ -33,9 +35,9 @@ embeddingsApp.registerTask({ name: "embeddings" }, async logger.info({ photoId, userId, requestId }, "Generating photo embeddings"); - const photo = await prisma.photo.findUnique({ + const photo = await prisma.media.findUnique({ where: { - photoId: { id: photoId, userId }, + mediaId: { id: photoId, userId, provider }, classification: { not: DbNull }, }, select: { @@ -76,7 +78,7 @@ embeddingsApp.registerTask({ name: "embeddings" }, async const imageVecStr = `[${(result.image ?? []).join(",")}]`; await prisma.$executeRaw( - Prisma.sql`UPDATE photos SET tag_vec = ${textVecStr}::vector, image_vec = ${imageVecStr}::vector WHERE id = ${photoId} AND user_id = ${userId}`, + Prisma.sql`UPDATE media SET tag_vec = ${textVecStr}::vector, image_vec = ${imageVecStr}::vector WHERE external_id = ${photoId} AND user_id = ${userId} AND provider = ${provider}`, ); logger.info({ photoId, userId, requestId }, "Photo embeddings generated"); diff --git a/apps/server/src/queue/image-collector.ts b/apps/server/src/queue/image-collector.ts index bfc2ed77..667af26c 100644 --- a/apps/server/src/queue/image-collector.ts +++ b/apps/server/src/queue/image-collector.ts @@ -1,13 +1,16 @@ import { Absurd } from "absurd-sdk"; import { env, prisma } from "@starlight/utils"; -import { http } from "@starlight/utils/http"; -import type { Tweet } from "@the-convocation/twitter-scraper"; -import UserAgent from "user-agents"; +import sharp from "sharp"; import { logger } from "@/logger"; import { absurdLogger, QUEUES, RETRY } from "@/queue/absurd"; import { classificationApp } from "@/queue/classification"; import { findDuplicatesByImageContent } from "@/services/duplicate-detection"; import { calculatePerceptualHash } from "@/services/image"; +import { + MAX_MEDIA_DOWNLOAD_BYTES, + MAX_POST_DOWNLOAD_BYTES, + readResponseBounded, +} from "@/services/media-download"; import { s3 } from "@/storage"; export const imagesApp = new Absurd({ @@ -16,161 +19,140 @@ export const imagesApp = new Absurd({ queueName: QUEUES.images, }); -export interface ImageCollectorJobData { - tweet: Tweet; - // From database +export interface MediaCollectorJobData { userId: string; + post: { + provider: string; + externalId: string; + sourceUrl: string; + authorExternalId?: string; + authorName?: string; + authorUsername?: string; + title?: string; + text?: string; + providerPayload: object; + media: Array<{ + externalId: string; + url: string; + kind?: string; + position: number; + fetchHeaders?: Record; + }>; + }; } -imagesApp.registerTask({ name: "images-collector" }, async (data) => { - const { tweet, userId } = data; - - // Tweet guaranteed to have IDs, fucking types - const id = tweet.id!; - - logger.info({ tweetId: tweet.id, userId }, "Processing tweet"); - - if (tweet.photos.length === 0) { - logger.debug({ tweetId: tweet.id, userId }, "Tweet has no photos, skipping job"); - return; - } - - const userAgent = new UserAgent(); - - // We can safely update Tweet record here, because we created Tweet object in scrapper queue - const tweetRecord = await prisma.tweet.update({ - where: { tweetId: { userId, id } }, - data: { - tweetData: tweet, +imagesApp.registerTask({ name: "images-collector" }, async (data) => { + const { post, userId } = data; + let downloadedBytes = 0; + const postRecord = await prisma.post.upsert({ + where: { postId: { id: post.externalId, userId, provider: post.provider } }, + create: { + id: post.externalId, + userId, + provider: post.provider, + sourceUrl: post.sourceUrl, + authorExternalId: post.authorExternalId, + authorName: post.authorName, + authorUsername: post.authorUsername, + title: post.title, + tweetData: { ...post.providerPayload, text: post.text, username: post.authorUsername }, photos: { createMany: { - data: tweet.photos.map((photo) => ({ - id: photo.id, - originalUrl: photo.url, + data: post.media.map((media) => ({ + id: media.externalId, + provider: post.provider, + position: media.position, + kind: media.kind ?? "image", + originalUrl: media.url, })), - // Guaranteed that if we'll restart a job then we won't have additional photos in Tweet relation skipDuplicates: true, }, }, }, - include: { - photos: true, + update: { + sourceUrl: post.sourceUrl, + authorExternalId: post.authorExternalId, + authorName: post.authorName, + authorUsername: post.authorUsername, + title: post.title, + tweetData: { ...post.providerPayload, text: post.text, username: post.authorUsername }, + photos: { + createMany: { + data: post.media.map((media) => ({ + id: media.externalId, + provider: post.provider, + position: media.position, + kind: media.kind ?? "image", + originalUrl: media.url, + })), + skipDuplicates: true, + }, + }, }, + include: { photos: true }, }); - logger.info( - { tweetId: tweet.id, userId, photos: tweetRecord.photos.length }, - "Tweet upserted with photos", - ); - - const refreshedPhotoIds = new Set(); - const refreshTimestamp = new Date(); - - for (const photo of tweetRecord.photos) { - if (photo.s3Path && photo.perceptualHash) { - logger.debug( - { - tweetId: tweet.id, - photoId: photo.id, - userId, - }, - "Photo already downloaded; skipping", - ); + for (const media of postRecord.photos) { + if (media.s3Path && (media.kind !== "image" || media.perceptualHash)) { continue; } - - const response = await http(photo.originalUrl, { - headers: { - "User-Agent": userAgent.toString(), - }, - }); - - if (!response.ok) { - logger.error( - { - tweetId: tweet.id, - photoUrl: photo.originalUrl, - status: response.status, - userId, - }, - "Failed to fetch photo", - ); - throw new Error(`Failed to fetch photo ${photo.originalUrl}`); + const input = post.media.find((item) => item.externalId === media.id); + if (!input) { + throw new Error(`Media ${media.id} is missing from collector payload`); + } + const remainingBytes = MAX_POST_DOWNLOAD_BYTES - downloadedBytes; + if (remainingBytes <= 0) { + throw new Error("Post media is too large"); + } + const bytes = await readResponseBounded( + await fetch(media.originalUrl, { headers: input.fetchHeaders }), + Math.min(MAX_MEDIA_DOWNLOAD_BYTES, remainingBytes), + ); + downloadedBytes += bytes.byteLength; + const extension = new URL(media.originalUrl).pathname.split(".").at(-1) ?? "jpg"; + const mediaPath = `media/${post.provider}/${userId}/${media.id}.${extension}`; + + if (media.kind !== "image") { + await s3.write(mediaPath, bytes); + await prisma.media.update({ + where: { mediaId: { id: media.id, userId, provider: post.provider } }, + data: { s3Path: mediaPath }, + }); + continue; } - const imageBuffer = await response.arrayBuffer(); - - const similarPhotos = await findDuplicatesByImageContent(imageBuffer); - - if (similarPhotos.length > 0) { - const existingPhoto = similarPhotos.find((similarPhoto) => similarPhoto.userId === userId); - - if (existingPhoto && !refreshedPhotoIds.has(existingPhoto.id)) { - await prisma.photo.update({ - where: { photoId: { id: existingPhoto.id, userId } }, - data: { updatedAt: refreshTimestamp }, - }); - refreshedPhotoIds.add(existingPhoto.id); - } - + const duplicates = await findDuplicatesByImageContent(bytes); + if (duplicates.length > 0) { logger.info( - { - tweetId: tweet.id, - photoId: photo.id, - userId, - refreshedPhotoId: existingPhoto?.id, - similarPhotos, - }, - "Found similar photos, skipping saving photo", + { mediaId: media.id, provider: post.provider, userId }, + "Duplicate media skipped", ); continue; } - - const extension = photo.originalUrl.split(".").pop() ?? "jpg"; - - const photoName = `${photo.externalId}.${extension}`; - const [, hash, metadata] = await Promise.all([ - s3.write(`media/${photoName}`, imageBuffer), - calculatePerceptualHash(imageBuffer), - new Bun.Image(imageBuffer).metadata().catch(() => ({ height: null, width: null })), + s3.write(mediaPath, bytes), + calculatePerceptualHash(bytes), + sharp(bytes) + .metadata() + .catch(() => ({ height: null, width: null })), ]); - - await prisma.photo.update({ - where: { photoId: { id: photo.id, userId } }, + await prisma.media.update({ + where: { mediaId: { id: media.id, userId, provider: post.provider } }, data: { perceptualHash: hash, - s3Path: `media/${photoName}`, + s3Path: mediaPath, height: metadata.height, width: metadata.width, }, }); - - // Enqueue classification job - try { - await classificationApp.spawn( - "classification", - { photoId: photo.id, userId }, - { - idempotencyKey: `classify-${photo.id}-${userId}`, - maxAttempts: 5, - retryStrategy: RETRY.classification, - }, - ); - } catch (error) { - logger.error( - { err: error, photoId: photo.id, userId }, - "Failed to enqueue classification job", - ); - } - - logger.info( + await classificationApp.spawn( + "classification", + { photoId: media.id, provider: post.provider, userId }, { - tweetId: tweet.id, - photoId: photo.id, - userId, + idempotencyKey: `classify-${post.provider}-${userId}-${media.id}`, + maxAttempts: 5, + retryStrategy: RETRY.classification, }, - "Photo saved to S3", ); } }); diff --git a/apps/server/src/queue/pixiv.ts b/apps/server/src/queue/pixiv.ts new file mode 100644 index 00000000..86303703 --- /dev/null +++ b/apps/server/src/queue/pixiv.ts @@ -0,0 +1,166 @@ +import { Absurd } from "absurd-sdk"; +import { withPixivClient } from "@starlight/api/services/pixiv-credential"; +import { env, prisma } from "@starlight/utils"; +import { absurdLogger, QUEUES, RETRY } from "@/queue/absurd"; +import { imagesApp, type MediaCollectorJobData } from "@/queue/image-collector"; + +const CONSECUTIVE_THRESHOLD = 15; +const SCHEDULE_INTERVAL_SECONDS = 60 * 60 * 6; + +export const pixivApp = new Absurd({ + db: env.DATABASE_URL, + log: absurdLogger, + queueName: QUEUES.pixiv, +}); + +export interface PixivCrawlJobData { + userId: string; + count: number; + limit: number; + cursor?: number; + visibility?: "public" | "private"; + force?: boolean; +} + +export interface ScheduledPixivJobData { + generation: number; + limit: number; + userId: string; +} + +export const getScheduledPixivGeneration = (date = new Date()) => + Math.floor(date.getTime() / (SCHEDULE_INTERVAL_SECONDS * 1000)); + +pixivApp.registerTask( + { name: "scheduled-pixiv-bookmarks" }, + async (data, ctx) => { + await ctx.sleepFor("next-run", SCHEDULE_INTERVAL_SECONDS); + const credential = await prisma.providerCredential.findUnique({ + where: { userId_provider: { userId: data.userId, provider: "pixiv" } }, + select: { provider: true }, + }); + if (credential) { + await pixivApp.spawn( + "pixiv-bookmarks", + { count: 0, limit: data.limit, userId: data.userId }, + { + idempotencyKey: `scheduled-pixiv-run-${data.userId}-${data.generation}`, + maxAttempts: 3, + retryStrategy: RETRY.pixiv, + }, + ); + } + const nextGeneration = data.generation + 1; + await pixivApp.spawn( + "scheduled-pixiv-bookmarks", + { ...data, generation: nextGeneration }, + { + idempotencyKey: `scheduled-pixiv-${data.userId}-${nextGeneration}`, + maxAttempts: 3, + retryStrategy: RETRY.pixiv, + }, + ); + }, +); + +pixivApp.registerTask({ name: "pixiv-bookmarks" }, async (data) => { + if (!data.visibility) { + const user = await prisma.user.findUnique({ + where: { id: data.userId }, + select: { pixivIncludePrivate: true, providerCredentials: { where: { provider: "pixiv" } } }, + }); + if (!user?.providerCredentials.length) { + return; + } + const visibilities: Array<"public" | "private"> = ["public"]; + if (user.pixivIncludePrivate) { + visibilities.push("private"); + } + await Promise.all( + visibilities.map((visibility) => + pixivApp.spawn( + "pixiv-bookmarks", + { ...data, count: 0, cursor: undefined, visibility }, + { + idempotencyKey: `pixiv-${data.userId}-${visibility}-start-${data.force ? "force" : "normal"}`, + maxAttempts: 3, + retryStrategy: RETRY.pixiv, + }, + ), + ), + ); + return; + } + + const page = await withPixivClient(data.userId, (client) => + client.bookmarks({ cursor: data.cursor, visibility: data.visibility! }), + ); + if (!page) { + return; + } + const known = new Set( + ( + await prisma.post.findMany({ + where: { + userId: data.userId, + provider: "pixiv", + id: { in: page.artworks.map((artwork) => artwork.id) }, + photos: { every: { s3Path: { not: null } } }, + }, + select: { id: true }, + }) + ).map((post) => post.id), + ); + + let consecutiveKnown = 0; + for (const artwork of page.artworks) { + consecutiveKnown = known.has(artwork.id) ? consecutiveKnown + 1 : 0; + const job: MediaCollectorJobData = { + userId: data.userId, + post: { + provider: "pixiv", + externalId: artwork.id, + sourceUrl: artwork.sourceUrl, + authorExternalId: artwork.author.id, + authorName: artwork.author.name, + authorUsername: artwork.author.username, + title: artwork.title, + text: artwork.caption, + providerPayload: { ...artwork.payload, starlightMediaType: artwork.type }, + media: artwork.mediaUrls.map((url, position) => ({ + externalId: `${artwork.id}:${position}`, + url, + position, + kind: "image", + fetchHeaders: { Referer: "https://www.pixiv.net/" }, + })), + }, + }; + await imagesApp.spawn("images-collector", job, { + idempotencyKey: `media-pixiv-${data.userId}-${artwork.id}`, + maxAttempts: 3, + retryStrategy: RETRY.images, + }); + if (!data.force && consecutiveKnown >= CONSECUTIVE_THRESHOLD) { + break; + } + } + + const count = data.count + page.artworks.length; + if ( + (!data.force && consecutiveKnown >= CONSECUTIVE_THRESHOLD) || + count >= data.limit || + !page.nextCursor + ) { + return; + } + await pixivApp.spawn( + "pixiv-bookmarks", + { ...data, count, cursor: page.nextCursor }, + { + idempotencyKey: `pixiv-${data.userId}-${data.visibility}-${page.nextCursor}`, + maxAttempts: 3, + retryStrategy: RETRY.pixiv, + }, + ); +}); diff --git a/apps/server/src/queue/scrapper.ts b/apps/server/src/queue/scrapper.ts index 7b09bf2f..82832a4a 100644 --- a/apps/server/src/queue/scrapper.ts +++ b/apps/server/src/queue/scrapper.ts @@ -2,11 +2,12 @@ import { Absurd } from "absurd-sdk"; import { CookieEncryption } from "@starlight/crypto"; import type { User } from "@starlight/utils"; import { env, prisma } from "@starlight/utils"; -import { type QueryTweetsResponse, Scraper, type Tweet } from "@the-convocation/twitter-scraper"; +import { type QueryTweetsResponse, Scraper } from "@the-convocation/twitter-scraper"; import { bot } from "@/bot"; import { logger } from "@/logger"; import { absurdLogger, QUEUES, RETRY } from "@/queue/absurd"; import { imagesApp } from "@/queue/image-collector"; +import type { MediaCollectorJobData } from "@/queue/image-collector"; import { Cookies } from "@/storage"; const cookieEncryption = new CookieEncryption( @@ -189,9 +190,10 @@ scrapperApp.registerTask({ name: "feed-scrapper" }, async (data const existingTweetMap = new Map( ( - await prisma.tweet.findMany({ + await prisma.post.findMany({ where: { userId, + provider: "twitter", id: { in: tweetIds }, photos: { every: { s3Path: { not: null } } }, }, @@ -201,13 +203,7 @@ scrapperApp.registerTask({ name: "feed-scrapper" }, async (data ); // Step 2: Process tweets and build batch operations - const newTweets: Array<{ - id: string; - userId: string; - tweetData: Tweet; - }> = []; - const updatedTweets: Array<{ id: string; tweetData: Tweet }> = []; - const tweetsToQueue: Array<{ tweet: Tweet; userId: string }> = []; + const tweetsToQueue: MediaCollectorJobData[] = []; let consecutiveKnownTweets = 0; let newTweetsInBatch = 0; @@ -222,22 +218,31 @@ scrapperApp.registerTask({ name: "feed-scrapper" }, async (data if (isNewTweet) { consecutiveKnownTweets = 0; newTweetsInBatch++; - newTweets.push({ - id: tweet.id, - userId, - tweetData: tweet, - }); } else { consecutiveKnownTweets++; - updatedTweets.push({ - id: tweet.id, - tweetData: tweet, - }); } // Only queue tweets with photos for image processing if (tweet.photos.length > 0) { - tweetsToQueue.push({ tweet, userId }); + tweetsToQueue.push({ + userId, + post: { + provider: "twitter", + externalId: tweet.id, + sourceUrl: `https://x.com/i/status/${tweet.id}`, + authorExternalId: tweet.userId, + authorName: tweet.name, + authorUsername: tweet.username, + text: tweet.text, + providerPayload: tweet, + media: tweet.photos.map((photo, position) => ({ + externalId: photo.id, + url: photo.url, + position, + kind: "image", + })), + }, + }); } // Stop if we've seen too many consecutive known tweets (unless force is enabled) @@ -255,35 +260,12 @@ scrapperApp.registerTask({ name: "feed-scrapper" }, async (data } } - // Step 3: Execute batch operations in transaction - await prisma.$transaction(async (tx) => { - // Batch create new tweets - if (newTweets.length > 0) { - await tx.tweet.createMany({ - data: newTweets, - skipDuplicates: true, - }); - } - - // Batch update existing tweets - if (updatedTweets.length > 0) { - await Promise.all( - updatedTweets.map((tweet) => - tx.tweet.update({ - where: { tweetId: { userId, id: tweet.id } }, - data: { tweetData: tweet.tweetData }, - }), - ), - ); - } - }); - // Queue image processing jobs for tweets with photos if (tweetsToQueue.length > 0) { await Promise.all( tweetsToQueue.map((job) => imagesApp.spawn("images-collector", job, { - idempotencyKey: `post-${job.tweet.id}-${job.userId}`, + idempotencyKey: `media-twitter-${job.userId}-${job.post.externalId}`, maxAttempts: 3, retryStrategy: RETRY.images, }), diff --git a/apps/server/src/scripts/classification.ts b/apps/server/src/scripts/classification.ts index a6826536..62ef9d2d 100644 --- a/apps/server/src/scripts/classification.ts +++ b/apps/server/src/scripts/classification.ts @@ -48,14 +48,14 @@ async function main() { } const photos = ALL_PICTURES - ? await prisma.photo.findMany({ + ? await prisma.media.findMany({ where: { deletedAt: null, s3Path: { not: null } }, - select: { id: true, userId: true }, + select: { id: true, provider: true, userId: true }, orderBy: { id: "asc" }, }) - : await prisma.$queryRaw<{ id: string; userId: string }[]>` - SELECT id, user_id as "userId" - FROM photos + : await prisma.$queryRaw<{ id: string; provider: string; userId: string }[]>` + SELECT external_id AS id, provider, user_id as "userId" + FROM media WHERE deleted_at IS NULL AND s3_path IS NOT NULL AND ( @@ -74,7 +74,7 @@ async function main() { if (!DRY_RUN && photos.length > 0) { await Promise.all( photos.map((photo) => { - const data = { photoId: photo.id, userId: photo.userId }; + const data = { photoId: photo.id, provider: photo.provider, userId: photo.userId }; return FORCE ? classificationApp.spawn("classification", data, { @@ -82,7 +82,7 @@ async function main() { retryStrategy: RETRY.classification, }) : classificationApp.spawn("classification", data, { - idempotencyKey: `classify-${data.photoId}-${data.userId}`, + idempotencyKey: `classify-${data.provider}-${data.photoId}-${data.userId}`, maxAttempts: 5, retryStrategy: RETRY.classification, }); diff --git a/apps/server/src/scripts/embeddings.ts b/apps/server/src/scripts/embeddings.ts index b1fd290b..975d51bc 100644 --- a/apps/server/src/scripts/embeddings.ts +++ b/apps/server/src/scripts/embeddings.ts @@ -43,16 +43,16 @@ async function main() { logger.info("Embeddings queue dropped and recreated"); } - const photos = await prisma.$queryRaw<{ id: string; userId: string }[]>` - SELECT id, user_id as "userId" - FROM photos + const photos = await prisma.$queryRaw<{ id: string; provider: string; userId: string }[]>` + SELECT external_id AS id, provider, user_id as "userId" + FROM media WHERE deleted_at IS NULL AND s3_path IS NOT NULL AND ( image_vec IS NULL OR tag_vec IS NULL ) - ORDER BY id ASC + ORDER BY external_id ASC `; let enqueued = 0; @@ -60,7 +60,7 @@ async function main() { if (!DRY_RUN && photos.length > 0) { await Promise.all( photos.map((photo) => { - const data = { photoId: photo.id, userId: photo.userId }; + const data = { photoId: photo.id, provider: photo.provider, userId: photo.userId }; return FORCE ? embeddingsApp.spawn("embeddings", data, { @@ -68,7 +68,7 @@ async function main() { retryStrategy: RETRY.embeddings, }) : embeddingsApp.spawn("embeddings", data, { - idempotencyKey: `embed-${data.photoId}-${data.userId}`, + idempotencyKey: `embed-${data.provider}-${data.photoId}-${data.userId}`, maxAttempts: 5, retryStrategy: RETRY.embeddings, }); diff --git a/apps/server/src/scripts/update-photo-dimensions.ts b/apps/server/src/scripts/update-photo-dimensions.ts index 4d1ee491..dbb2a4ec 100644 --- a/apps/server/src/scripts/update-photo-dimensions.ts +++ b/apps/server/src/scripts/update-photo-dimensions.ts @@ -21,7 +21,7 @@ async function main() { ); // Find photos with null height or width that have s3Path - const photos = await prisma.photo.findMany({ + const photos = await prisma.media.findMany({ where: { deletedAt: null, s3Path: { not: null }, @@ -30,6 +30,7 @@ async function main() { select: { id: true, userId: true, + provider: true, s3Path: true, height: true, width: true, @@ -85,8 +86,14 @@ async function main() { if (!DRY_RUN) { // Update photo with dimensions - await prisma.photo.update({ - where: { photoId: { id: photo.id, userId: photo.userId } }, + await prisma.media.update({ + where: { + mediaId: { + id: photo.id, + userId: photo.userId, + provider: photo.provider, + }, + }, data: { height: metadata.height, width: metadata.width, diff --git a/apps/server/src/services/duplicate-detection.ts b/apps/server/src/services/duplicate-detection.ts index d89891f5..23c3d75f 100644 --- a/apps/server/src/services/duplicate-detection.ts +++ b/apps/server/src/services/duplicate-detection.ts @@ -8,6 +8,7 @@ interface SimilarPhoto { originalUrl: string; perceptualHash: string; s3Path?: string; + sourceUrl: string; tweetId: string; userId: string; } @@ -29,7 +30,7 @@ export async function findSimilarPhotos( logger.debug({ prefix, field, maxCandidates }, "Searching for similar photos"); - const candidates = await prisma.photo.findMany({ + const candidates = await prisma.media.findMany({ where: { [field]: prefix, perceptualHash: { not: null }, @@ -48,6 +49,7 @@ export async function findSimilarPhotos( s3Path: true, originalUrl: true, tweetId: true, + tweet: { select: { sourceUrl: true } }, }, take: maxCandidates, }); @@ -72,6 +74,7 @@ export async function findSimilarPhotos( s3Path: candidate.s3Path || undefined, originalUrl: candidate.originalUrl, tweetId: candidate.tweetId, + sourceUrl: candidate.tweet.sourceUrl, }); } } diff --git a/apps/server/src/services/media-download.ts b/apps/server/src/services/media-download.ts new file mode 100644 index 00000000..3acf701a --- /dev/null +++ b/apps/server/src/services/media-download.ts @@ -0,0 +1,42 @@ +export const MAX_MEDIA_DOWNLOAD_BYTES = 50_000_000; +export const MAX_POST_DOWNLOAD_BYTES = 200_000_000; + +export const readResponseBounded = async (response: Response, limit = MAX_MEDIA_DOWNLOAD_BYTES) => { + if (!response.ok) { + throw new Error(`Media request failed (${response.status})`); + } + const declaredLength = Number(response.headers.get("content-length")); + if (Number.isFinite(declaredLength) && declaredLength > limit) { + await response.body?.cancel(); + throw new Error("Media is too large"); + } + if (!response.body) { + throw new Error("Media response had no body"); + } + const reader = response.body.getReader(); + const chunks: Uint8Array[] = []; + let total = 0; + try { + while (true) { + const chunk = await reader.read(); + if (chunk.done) { + break; + } + total += chunk.value.byteLength; + if (total > limit) { + await reader.cancel("Media is too large"); + throw new Error("Media is too large"); + } + chunks.push(chunk.value); + } + } finally { + reader.releaseLock(); + } + const output = new Uint8Array(total); + let offset = 0; + for (const chunk of chunks) { + output.set(chunk, offset); + offset += chunk.byteLength; + } + return output; +}; diff --git a/apps/server/src/services/pixiv-media.test.ts b/apps/server/src/services/pixiv-media.test.ts new file mode 100644 index 00000000..9bb6e9c2 --- /dev/null +++ b/apps/server/src/services/pixiv-media.test.ts @@ -0,0 +1,103 @@ +import { describe, expect, test } from "bun:test"; +import { readdir } from "node:fs/promises"; +import { tmpdir } from "node:os"; +import JSZip from "jszip"; +import { readResponseBounded } from "./media-download"; +import { + buildFfmpegConcat, + extractUgoiraZip, + parsePixivArtworkUrl, +} from "./pixiv-media"; + +describe("parsePixivArtworkUrl", () => { + test("accepts only canonical HTTPS artwork URLs", () => { + expect(parsePixivArtworkUrl("https://www.pixiv.net/artworks/12345")).toBe( + "12345" + ); + expect(parsePixivArtworkUrl("https://pixiv.net/artworks/12345/")).toBe( + "12345" + ); + expect( + parsePixivArtworkUrl("http://www.pixiv.net/artworks/12345") + ).toBeNull(); + expect( + parsePixivArtworkUrl("https://pixiv.net.evil.test/artworks/12345") + ).toBeNull(); + expect( + parsePixivArtworkUrl("https://www.pixiv.net/users/12345") + ).toBeNull(); + }); +}); + +describe("ugoira conversion input", () => { + test("preserves variable frame delays", () => { + expect( + buildFfmpegConcat([ + { file: "a.jpg", delay: 40 }, + { file: "b.jpg", delay: 125 }, + ]) + ).toContain("duration 0.04\nfile 'b.jpg'\nduration 0.125"); + }); + + test("rejects traversal in frame metadata", async () => { + const zip = new JSZip(); + zip.file("frame.jpg", "data"); + const archive = await zip.generateAsync({ type: "arraybuffer" }); + expect( + extractUgoiraZip(archive, [{ file: "../frame.jpg", delay: 100 }]) + ).rejects.toThrow("Unsafe ugoira archive path"); + }); + + test("bounds declared and streamed download sizes", async () => { + const declared = new Response("small", { + headers: { "content-length": "100" }, + }); + await expect(readResponseBounded(declared, 10)).rejects.toThrow( + "too large" + ); + const streamed = new Response( + new ReadableStream({ + start(controller) { + controller.enqueue(new Uint8Array(8)); + controller.enqueue(new Uint8Array(8)); + controller.close(); + }, + }) + ); + await expect(readResponseBounded(streamed, 10)).rejects.toThrow( + "too large" + ); + }); + + test("bounds cumulative extraction and removes failed temp directories", async () => { + const before = new Set( + (await readdir(tmpdir())).filter((name) => + name.startsWith("starlight-ugoira-") + ) + ); + const zip = new JSZip(); + zip.file("a.jpg", new Uint8Array(8)); + zip.file("b.jpg", new Uint8Array(8)); + const archive = await zip.generateAsync({ type: "arraybuffer" }); + expect( + extractUgoiraZip( + archive, + [ + { file: "a.jpg", delay: 100 }, + { file: "b.jpg", delay: 100 }, + ], + { uncompressed: 10 } + ) + ).rejects.toThrow("too large"); + const after = (await readdir(tmpdir())).filter((name) => + name.startsWith("starlight-ugoira-") + ); + const leakedDirectories: string[] = []; + for (const entry of after) { + if (!before.has(entry)) { + leakedDirectories.push(entry); + } + } + expect(leakedDirectories).toEqual([]); + }); +}); diff --git a/apps/server/src/services/pixiv-media.ts b/apps/server/src/services/pixiv-media.ts new file mode 100644 index 00000000..fdab01ad --- /dev/null +++ b/apps/server/src/services/pixiv-media.ts @@ -0,0 +1,159 @@ +import { createWriteStream } from "node:fs"; +import { mkdir, mkdtemp, rm, stat, writeFile } from "node:fs/promises"; +import { tmpdir } from "node:os"; +import { basename, join } from "node:path"; +import { Transform } from "node:stream"; +import { pipeline } from "node:stream/promises"; +import JSZip from "jszip"; +import { MAX_MEDIA_DOWNLOAD_BYTES } from "@/services/media-download"; + +const PIXIV_HOSTS = new Set(["pixiv.net", "www.pixiv.net"]); +export const MAX_PIXIV_DOWNLOAD_BYTES = MAX_MEDIA_DOWNLOAD_BYTES; +export const MAX_UGOIRA_UNCOMPRESSED_BYTES = 200_000_000; +const MAX_UGOIRA_FRAMES = 1000; +const FFMPEG_TIMEOUT_MS = 120_000; + +export const parsePixivArtworkUrl = (value: string) => { + let url: URL; + try { + url = new URL(value); + } catch { + return null; + } + if (url.protocol !== "https:" || !PIXIV_HOSTS.has(url.hostname.toLowerCase())) { + return null; + } + const match = /^\/artworks\/(\d+)\/?$/.exec(url.pathname); + return match?.[1] ?? null; +}; + +export const buildFfmpegConcat = (frames: Array<{ file: string; delay: number }>) => { + if (frames.length === 0 || frames.some((frame) => frame.delay <= 0)) { + throw new Error("Invalid ugoira frame timing"); + } + const lines: string[] = []; + for (const frame of frames) { + if (frame.file.includes("'") || frame.file.includes("\n")) { + throw new Error("Invalid ugoira frame filename"); + } + lines.push(`file '${frame.file}'`, `duration ${frame.delay / 1000}`); + } + lines.push(`file '${frames.at(-1)?.file}'`); + return `${lines.join("\n")}\n`; +}; + +export const extractUgoiraZip = async ( + archive: ArrayBuffer | Uint8Array, + frames: Array<{ file: string; delay: number }>, + limits: { compressed?: number; uncompressed?: number; frames?: number } = {}, +) => { + const compressedLimit = limits.compressed ?? MAX_PIXIV_DOWNLOAD_BYTES; + const uncompressedLimit = limits.uncompressed ?? MAX_UGOIRA_UNCOMPRESSED_BYTES; + const frameLimit = limits.frames ?? MAX_UGOIRA_FRAMES; + if (archive.byteLength > compressedLimit || frames.length > frameLimit) { + throw new Error("Ugoira is too large"); + } + for (const frame of frames) { + if (basename(frame.file) !== frame.file || frame.file.includes("\\")) { + throw new Error("Unsafe ugoira archive path"); + } + } + const zip = await JSZip.loadAsync(archive); + for (const entry of Object.values(zip.files)) { + if (entry.unsafeOriginalName && entry.unsafeOriginalName !== entry.name) { + throw new Error("Unsafe ugoira archive path"); + } + } + const directory = await mkdtemp(join(tmpdir(), "starlight-ugoira-")); + let total = 0; + try { + await mkdir(directory, { recursive: true }); + for (const frame of frames) { + const entry = zip.file(frame.file); + if (!entry || entry.dir) { + throw new Error(`Missing ugoira frame ${frame.file}`); + } + const limiter = new Transform({ + transform(chunk: Buffer, _encoding, callback) { + total += chunk.byteLength; + if (total > uncompressedLimit) { + callback(new Error("Ugoira is too large")); + return; + } + callback(null, chunk); + }, + }); + await pipeline( + entry.nodeStream("nodebuffer"), + limiter, + createWriteStream(join(directory, frame.file), { flags: "wx" }), + ); + } + const concatPath = join(directory, "frames.txt"); + await writeFile( + concatPath, + buildFfmpegConcat(frames.map((frame) => ({ ...frame, file: join(directory, frame.file) }))), + ); + return { directory, concatPath }; + } catch (error) { + await rm(directory, { recursive: true, force: true }); + throw error; + } +}; + +export const convertUgoira = async (concatPath: string, output: string) => { + // biome-ignore lint/correctness/noUndeclaredVariables: Server runtime is Bun. + const child = Bun.spawn( + [ + "ffmpeg", + "-y", + "-f", + "concat", + "-safe", + "0", + "-i", + concatPath, + "-vf", + "scale=trunc(iw/2)*2:trunc(ih/2)*2", + "-c:v", + "libx264", + "-pix_fmt", + "yuv420p", + "-movflags", + "+faststart", + "-fs", + String(MAX_PIXIV_DOWNLOAD_BYTES), + output, + ], + { stdout: "ignore", stderr: "pipe" }, + ); + const stderr = new Response(child.stderr).arrayBuffer(); + let timeout: ReturnType | undefined; + try { + const result = await Promise.race([ + child.exited, + new Promise<"timeout">((resolve) => { + timeout = setTimeout(() => { + resolve("timeout"); + }, FFMPEG_TIMEOUT_MS); + }), + ]); + if (result === "timeout") { + child.kill(); + await child.exited; + await stderr; + throw new Error("Ugoira conversion timed out"); + } + await stderr; + if (result !== 0) { + throw new Error("Failed to convert ugoira"); + } + if ((await stat(output)).size > MAX_PIXIV_DOWNLOAD_BYTES) { + throw new Error("Converted ugoira is too large"); + } + } finally { + if (timeout) { + clearTimeout(timeout); + } + } +}; diff --git a/apps/web/src/routes/settings.tsx b/apps/web/src/routes/settings.tsx index d9372683..c9550a53 100644 --- a/apps/web/src/routes/settings.tsx +++ b/apps/web/src/routes/settings.tsx @@ -1,7 +1,7 @@ import type { ProfileResult } from "@starlight/api/routers/index"; import { useMutation, useQuery, useQueryClient } from "@tanstack/react-query"; import { createFileRoute, Link } from "@tanstack/react-router"; -import { AlertCircle, Cookie, Trash2 } from "lucide-react"; +import { AlertCircle, Cookie, KeyRound, Trash2 } from "lucide-react"; import { useState } from "react"; import { Alert, AlertDescription, AlertTitle } from "@/components/ui/alert"; import { Button } from "@/components/ui/button"; @@ -32,6 +32,7 @@ export const Route = createFileRoute("/settings")({ function RouteComponent() { const [newCookies, setNewCookies] = useState(""); + const [pixivToken, setPixivToken] = useState(""); const [displayError, setDisplayError] = useState(null); const { rawInitData } = useTelegramContext(); @@ -93,6 +94,38 @@ function RouteComponent() { }), ); + const savePixivMutation = useMutation( + orpc.pixiv.save.mutationOptions({ + onSuccess: () => { + queryClient.setQueryData(["profile"], (old: ProfileResult) => ({ + ...old, + hasPixivCredential: true, + })); + setPixivToken(""); + }, + }), + ); + const deletePixivMutation = useMutation( + orpc.pixiv.delete.mutationOptions({ + onSuccess: () => { + queryClient.setQueryData(["profile"], (old: ProfileResult) => ({ + ...old, + hasPixivCredential: false, + })); + }, + }), + ); + const pixivPrivateMutation = useMutation( + orpc.pixiv.privateBookmarks.mutationOptions({ + onSuccess: (_data, variables) => { + queryClient.setQueryData(["profile"], (old: ProfileResult) => ({ + ...old, + pixivIncludePrivate: variables.enabled, + })); + }, + }), + ); + if (isLoading && !profile) { return (
@@ -148,7 +181,10 @@ function RouteComponent() { const isSubmitting = saveCookiesMutation.isPending || deleteCookiesMutation.isPending || - visibilityMutation.isPending; + visibilityMutation.isPending || + savePixivMutation.isPending || + deletePixivMutation.isPending || + pixivPrivateMutation.isPending; return (
@@ -242,6 +278,61 @@ function RouteComponent() { )} +
+

+ Pixiv +

+ {profile?.hasPixivCredential ? ( + <> + + + Pixiv is connected. + + + + + ) : ( +
{ + event.preventDefault(); + savePixivMutation.mutate({ refreshToken: pixivToken }); + }} + > + + + + )} +
+ {/* Profile Visibility Section */}