Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
98 changes: 98 additions & 0 deletions apps/api/src/device-agent/device-agent.controller.spec.ts
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,8 @@ import { PermissionGuard } from '../auth/permission.guard';
import type { AuthContext as AuthContextType } from '../auth/types';
import { Readable } from 'stream';

jest.mock('@db', () => ({ db: {} }));

jest.mock('../auth/auth.server', () => ({
auth: { api: { getSession: jest.fn() } },
}));
Expand All @@ -26,6 +28,8 @@ describe('DeviceAgentController', () => {
const mockService = {
downloadMacAgent: jest.fn(),
downloadWindowsAgent: jest.fn(),
getUpdateFile: jest.fn(),
headUpdateFile: jest.fn(),
};

const mockGuard = { canActivate: jest.fn().mockReturnValue(true) };
Expand All @@ -42,6 +46,7 @@ describe('DeviceAgentController', () => {

const mockRes = {
set: jest.fn(),
redirect: jest.fn(),
};

beforeEach(async () => {
Expand Down Expand Up @@ -155,4 +160,97 @@ describe('DeviceAgentController', () => {
).rejects.toThrow('Agent not found');
});
});

describe('getUpdateFile', () => {
it('streams a yml manifest with cache headers', async () => {
const mockStream = Readable.from(Buffer.from('version: 1.0.5'));
mockService.getUpdateFile.mockResolvedValue({
kind: 'stream',
stream: mockStream,
contentType: 'text/yaml',
contentLength: 14,
});

const result = await controller.getUpdateFile(
'latest-mac.yml',
mockRes as never,
);

expect(result).toBeInstanceOf(StreamableFile);
expect(mockRes.set).toHaveBeenCalledWith(
expect.objectContaining({
'Content-Type': 'text/yaml',
'Cache-Control': 'public, max-age=300',
'Content-Length': '14',
}),
);
expect(mockRes.redirect).not.toHaveBeenCalled();
});

it('issues a 302 redirect to the presigned URL for binaries', async () => {
mockService.getUpdateFile.mockResolvedValue({
kind: 'redirect',
url: 'https://s3.example.com/signed-zip',
});

const result = await controller.getUpdateFile(
'CompAI-Device-Agent-1.0.5-arm64.zip',
mockRes as never,
);

expect(mockRes.redirect).toHaveBeenCalledWith(
302,
'https://s3.example.com/signed-zip',
);
expect(mockRes.set).toHaveBeenCalledWith(
expect.objectContaining({
'Cache-Control': 'no-store',
}),
);
expect(result).toBeUndefined();
});
});

describe('headUpdateFile', () => {
it('returns metadata headers for a yml manifest', async () => {
mockService.headUpdateFile.mockResolvedValue({
kind: 'stream',
contentType: 'text/yaml',
contentLength: 859,
});

const result = await controller.headUpdateFile(
'latest-mac.yml',
mockRes as never,
);

expect(mockRes.set).toHaveBeenCalledWith(
expect.objectContaining({
'Content-Type': 'text/yaml',
'Cache-Control': 'public, max-age=300',
'Content-Length': '859',
}),
);
expect(result).toBe('');
expect(mockRes.redirect).not.toHaveBeenCalled();
});

it('issues a 302 redirect for HEAD on binaries', async () => {
mockService.headUpdateFile.mockResolvedValue({
kind: 'redirect',
url: 'https://s3.example.com/signed-head',
});

const result = await controller.headUpdateFile(
'CompAI-Device-Agent-1.0.5-arm64.zip',
mockRes as never,
);

expect(mockRes.redirect).toHaveBeenCalledWith(
302,
'https://s3.example.com/signed-head',
);
expect(result).toBeUndefined();
});
});
});
12 changes: 12 additions & 0 deletions apps/api/src/device-agent/device-agent.controller.ts
Original file line number Diff line number Diff line change
Expand Up @@ -70,6 +70,12 @@ export class DeviceAgentController {
) {
const result = await this.deviceAgentService.getUpdateFile({ filename });

if (result.kind === 'redirect') {
res.set({ 'Cache-Control': 'no-store' });
res.redirect(302, result.url);
return;
}

res.set({
'Content-Type': result.contentType,
'Cache-Control': 'public, max-age=300',
Expand All @@ -90,6 +96,12 @@ export class DeviceAgentController {
) {
const result = await this.deviceAgentService.headUpdateFile({ filename });

if (result.kind === 'redirect') {
res.set({ 'Cache-Control': 'no-store' });
res.redirect(302, result.url);
return;
}

res.set({
'Content-Type': result.contentType,
'Cache-Control': 'public, max-age=300',
Expand Down
137 changes: 136 additions & 1 deletion apps/api/src/device-agent/device-agent.service.spec.ts
Original file line number Diff line number Diff line change
Expand Up @@ -5,10 +5,28 @@ import {
import { Readable } from 'stream';

const mockSend = jest.fn();
const mockGetSignedUrl = jest.fn();

class MockGetObjectCommand {
constructor(public readonly input: unknown) {
Object.assign(this, input as object);
}
}

class MockHeadObjectCommand {
constructor(public readonly input: unknown) {
Object.assign(this, input as object);
}
}

jest.mock('@aws-sdk/client-s3', () => ({
S3Client: jest.fn().mockImplementation(() => ({ send: mockSend })),
GetObjectCommand: jest.fn().mockImplementation((input) => input),
GetObjectCommand: MockGetObjectCommand,
HeadObjectCommand: MockHeadObjectCommand,
}));

jest.mock('@aws-sdk/s3-request-presigner', () => ({
getSignedUrl: (...args: unknown[]) => mockGetSignedUrl(...args),
}));

import { DeviceAgentService } from './device-agent.service';
Expand Down Expand Up @@ -92,6 +110,123 @@ describe('DeviceAgentService', () => {
});
});

describe('getUpdateFile', () => {
it('streams .yml manifests directly from S3', async () => {
const mockStream = new Readable({ read() {} });
mockSend.mockResolvedValue({
Body: mockStream,
ContentLength: 859,
});

const result = await service.getUpdateFile({ filename: 'latest-mac.yml' });

expect(result).toEqual({
kind: 'stream',
stream: mockStream,
contentType: 'text/yaml',
contentLength: 859,
});
expect(mockGetSignedUrl).not.toHaveBeenCalled();
});

it('redirects binary downloads to a presigned S3 URL signed for GET', async () => {
mockGetSignedUrl.mockResolvedValue('https://s3.example.com/signed-zip-url');

const result = await service.getUpdateFile({
filename: 'CompAI-Device-Agent-1.0.5-arm64.zip',
});

expect(result).toEqual({
kind: 'redirect',
url: 'https://s3.example.com/signed-zip-url',
});
expect(mockSend).not.toHaveBeenCalled();
expect(mockGetSignedUrl).toHaveBeenCalledTimes(1);
const [, command] = mockGetSignedUrl.mock.calls[0];
expect(command).toBeInstanceOf(MockGetObjectCommand);
expect(command).toMatchObject({
Bucket: 'test-bucket',
Key: 'device-agent/production/updates/CompAI-Device-Agent-1.0.5-arm64.zip',
});
});

it.each([
'CompAI-Device-Agent-1.0.5-arm64.zip',
'CompAI-Device-Agent-1.0.5-setup.exe',
'CompAI-Device-Agent-1.0.5-arm64.dmg',
'CompAI-Device-Agent-1.0.5-x86_64.AppImage',
'CompAI-Device-Agent-1.0.5-arm64.zip.blockmap',
])('redirects binary file %s', async (filename) => {
mockGetSignedUrl.mockResolvedValue('https://s3.example.com/signed');

const result = await service.getUpdateFile({ filename });

expect(result).toEqual({
kind: 'redirect',
url: 'https://s3.example.com/signed',
});
});

it('throws NotFoundException for invalid filenames', async () => {
await expect(
service.getUpdateFile({ filename: '../etc/passwd' }),
).rejects.toThrow(NotFoundException);
await expect(
service.getUpdateFile({ filename: 'foo.txt' }),
).rejects.toThrow(NotFoundException);
});

it('throws NotFoundException when S3 returns NoSuchKey for a yml manifest', async () => {
const error = new Error('Not found');
error.name = 'NoSuchKey';
mockSend.mockRejectedValue(error);

await expect(
service.getUpdateFile({ filename: 'latest-mac.yml' }),
).rejects.toThrow(NotFoundException);
});
});

describe('headUpdateFile', () => {
it('returns metadata for .yml manifests', async () => {
mockSend.mockResolvedValue({ ContentLength: 859 });

const result = await service.headUpdateFile({
filename: 'latest-mac.yml',
});

expect(result).toEqual({
kind: 'stream',
contentType: 'text/yaml',
contentLength: 859,
});
expect(mockGetSignedUrl).not.toHaveBeenCalled();
});

it('redirects binary HEAD requests to a URL signed with HeadObjectCommand', async () => {
mockGetSignedUrl.mockResolvedValue('https://s3.example.com/signed-head');

const result = await service.headUpdateFile({
filename: 'CompAI-Device-Agent-1.0.5-arm64.zip',
});

expect(result).toEqual({
kind: 'redirect',
url: 'https://s3.example.com/signed-head',
});
expect(mockSend).not.toHaveBeenCalled();
expect(mockGetSignedUrl).toHaveBeenCalledTimes(1);
const [, command] = mockGetSignedUrl.mock.calls[0];
// S3 signs each HTTP method separately; a GET-signed URL would be
// rejected when used with a HEAD request.
expect(command).toBeInstanceOf(MockHeadObjectCommand);
expect(command).toMatchObject({
Bucket: 'test-bucket',
Key: 'device-agent/production/updates/CompAI-Device-Agent-1.0.5-arm64.zip',
});
});
});

describe('downloadWindowsAgent', () => {
it('should return stream, filename, and contentType on success', async () => {
const mockStream = new Readable({ read() {} });
Expand Down
Loading
Loading