Refactor SSE Endpoint with Enhanced Token Validation and Client Management
- Replaced direct token comparison with TokenManager validation - Implemented robust SSE client connection and event subscription workflow - Added detailed client authentication and connection status reporting - Improved SSE endpoint with flexible event, entity, and domain subscription - Enhanced error handling and client disconnect management
This commit is contained in:
68
src/index.ts
68
src/index.ts
@@ -5,7 +5,7 @@ import { v4 as uuidv4 } from 'uuid';
|
|||||||
import { sseManager } from './sse/index.js';
|
import { sseManager } from './sse/index.js';
|
||||||
import { ILogger } from "@digital-alchemy/core";
|
import { ILogger } from "@digital-alchemy/core";
|
||||||
import express from 'express';
|
import express from 'express';
|
||||||
import { rateLimiter, securityHeaders, validateRequest, sanitizeInput, errorHandler } from './security/index.js';
|
import { rateLimiter, securityHeaders, validateRequest, sanitizeInput, errorHandler, TokenManager } from './security/index.js';
|
||||||
import { MCP_SCHEMA } from './mcp/schema.js';
|
import { MCP_SCHEMA } from './mcp/schema.js';
|
||||||
|
|
||||||
// Load environment variables based on NODE_ENV
|
// Load environment variables based on NODE_ENV
|
||||||
@@ -139,28 +139,70 @@ app.get('/subscribe_events', (req, res) => {
|
|||||||
// Get token from query parameter
|
// Get token from query parameter
|
||||||
const token = req.query.token?.toString();
|
const token = req.query.token?.toString();
|
||||||
|
|
||||||
if (!token || token !== HASS_TOKEN) {
|
if (!token || !TokenManager.validateToken(token)) {
|
||||||
return res.status(401).json({
|
return res.status(401).json({
|
||||||
success: false,
|
success: false,
|
||||||
message: 'Unauthorized - Invalid token'
|
message: 'Unauthorized - Invalid token'
|
||||||
});
|
});
|
||||||
}
|
}
|
||||||
|
|
||||||
const tool = tools.find(t => t.name === 'subscribe_events');
|
// Set SSE headers
|
||||||
if (!tool) {
|
res.writeHead(200, {
|
||||||
return res.status(404).json({
|
'Content-Type': 'text/event-stream',
|
||||||
success: false,
|
'Cache-Control': 'no-cache',
|
||||||
message: 'Tool not found'
|
'Connection': 'keep-alive',
|
||||||
|
'Access-Control-Allow-Origin': '*'
|
||||||
});
|
});
|
||||||
|
|
||||||
|
// Send initial connection message
|
||||||
|
res.write(`data: ${JSON.stringify({
|
||||||
|
type: 'connection',
|
||||||
|
status: 'connected',
|
||||||
|
timestamp: new Date().toISOString()
|
||||||
|
})}\n\n`);
|
||||||
|
|
||||||
|
const clientId = uuidv4();
|
||||||
|
const client = {
|
||||||
|
id: clientId,
|
||||||
|
send: (data: string) => {
|
||||||
|
res.write(`data: ${data}\n\n`);
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
// Add client to SSE manager
|
||||||
|
const sseClient = sseManager.addClient(client, token);
|
||||||
|
if (!sseClient || !sseClient.authenticated) {
|
||||||
|
res.write(`data: ${JSON.stringify({
|
||||||
|
type: 'error',
|
||||||
|
message: sseClient ? 'Authentication failed' : 'Maximum client limit reached',
|
||||||
|
timestamp: new Date().toISOString()
|
||||||
|
})}\n\n`);
|
||||||
|
return res.end();
|
||||||
}
|
}
|
||||||
|
|
||||||
tool.execute({
|
// Subscribe to events if specified
|
||||||
token,
|
const events = req.query.events?.toString().split(',').filter(Boolean);
|
||||||
events: req.query.events?.toString().split(','),
|
if (events?.length) {
|
||||||
entity_id: req.query.entity_id?.toString(),
|
events.forEach(event => sseManager.subscribeToEvent(clientId, event));
|
||||||
domain: req.query.domain?.toString(),
|
}
|
||||||
response: res
|
|
||||||
|
// Subscribe to entity if specified
|
||||||
|
const entityId = req.query.entity_id?.toString();
|
||||||
|
if (entityId) {
|
||||||
|
sseManager.subscribeToEntity(clientId, entityId);
|
||||||
|
}
|
||||||
|
|
||||||
|
// Subscribe to domain if specified
|
||||||
|
const domain = req.query.domain?.toString();
|
||||||
|
if (domain) {
|
||||||
|
sseManager.subscribeToDomain(clientId, domain);
|
||||||
|
}
|
||||||
|
|
||||||
|
// Handle client disconnect
|
||||||
|
req.on('close', () => {
|
||||||
|
sseManager.removeClient(clientId);
|
||||||
});
|
});
|
||||||
|
|
||||||
} catch (error) {
|
} catch (error) {
|
||||||
res.status(500).json({
|
res.status(500).json({
|
||||||
success: false,
|
success: false,
|
||||||
|
|||||||
@@ -1,5 +1,6 @@
|
|||||||
import { EventEmitter } from 'events';
|
import { EventEmitter } from 'events';
|
||||||
import { HassEntity, HassEvent, StateChangedEvent } from '../types/hass.js';
|
import { HassEntity, HassEvent, StateChangedEvent } from '../types/hass.js';
|
||||||
|
import { TokenManager } from '../security/index.js';
|
||||||
|
|
||||||
interface RateLimit {
|
interface RateLimit {
|
||||||
count: number;
|
count: number;
|
||||||
@@ -264,8 +265,7 @@ export class SSEManager extends EventEmitter {
|
|||||||
|
|
||||||
private validateToken(token?: string): boolean {
|
private validateToken(token?: string): boolean {
|
||||||
if (!token) return false;
|
if (!token) return false;
|
||||||
// Compare with HASS_TOKEN from environment
|
return TokenManager.validateToken(token);
|
||||||
return token === process.env.HASS_TOKEN;
|
|
||||||
}
|
}
|
||||||
|
|
||||||
// Utility methods
|
// Utility methods
|
||||||
|
|||||||
Reference in New Issue
Block a user