1 files changed, 87 insertions, 101 deletions

diff --git a/src/server/main.ts b/src/server/main.ts
index e94fcb0..f9e43d5 100644
--- a/src/server/main.ts
+++ b/src/server/main.ts
@@ -2,17 +2,16 @@ import dotenv from 'dotenv';
 import express, { Request, Response } from 'express';
 import { randomBytes } from 'crypto';
 import qrcode from 'qrcode';
-// TODO: Import digiidTs library once linked
-// import * as digiidTs from 'digiid-ts';
+// Import actual functions from the linked library
+import { generateDigiIDUri, verifyDigiIDCallback, DigiIDError } from 'digiid-ts';
 
 // Load environment variables from .env file
 dotenv.config();
 
 const app = express();
-const PORT = process.env.PORT || 3001; // Default to 3001 if PORT is not set
+const PORT = process.env.PORT || 3001;
 
 // In-memory storage for demo purposes
-// NOTE: In a production application, use a persistent store (e.g., Redis, database)
 interface SessionState {
   nonce: string;
   status: 'pending' | 'success' | 'failed';
@@ -23,14 +22,12 @@ const sessionStore = new Map<string, SessionState>();
 const nonceToSessionMap = new Map<string, string>();
 
 // Middleware
-app.use(express.json()); // Parse JSON bodies
+app.use(express.json());
 
 console.log('Server starting...');
 console.log(`Attempting to listen on port: ${PORT}`);
 console.log(`Configured PUBLIC_URL: ${process.env.PUBLIC_URL}`);
 
-// TODO: Implement DigiID API endpoints
-
 // Endpoint to initiate the DigiID authentication flow
 app.get('/api/digiid/start', async (req: Request, res: Response) => {
   try {
@@ -52,27 +49,25 @@ app.get('/api/digiid/start', async (req: Request, res: Response) => {
        return res.status(500).json({ error: 'Server configuration error: Invalid PUBLIC_URL format.' });
     }
 
-    // Determine if the callback URL is insecure (HTTP)
     const unsecure = callbackUrl.startsWith('http://');
-
-    // Placeholder for digiidTs.generateDigiIDUri call
-    // const digiIdUri = digiidTs.generateDigiIDUri({ callbackUrl, unsecure, nonce });
-    const digiIdUri = `digiid://example.com?x=${nonce}&unsecure=${unsecure ? 1 : 0}&callback=${encodeURIComponent(callbackUrl)}`; // TEMPORARY PLACEHOLDER
+    const digiIdUri = generateDigiIDUri({ callbackUrl, unsecure, nonce });
     console.log(`Generated DigiID URI: ${digiIdUri} for session ${sessionId}`);
 
-    // Store session state
     sessionStore.set(sessionId, { nonce, status: 'pending' });
     nonceToSessionMap.set(nonce, sessionId);
     console.log(`Stored pending session: ${sessionId}, nonce: ${nonce}`);
 
-    // Generate QR code
     const qrCodeDataUrl = await qrcode.toDataURL(digiIdUri);
-
     res.json({ sessionId, qrCodeDataUrl });
 
   } catch (error) {
     console.error('Error in /api/digiid/start:', error);
-    res.status(500).json({ error: 'Internal server error' });
+    // Check if it's a DigiIDError specifically from generateDigiIDUri
+    if (error instanceof DigiIDError) {
+        res.status(400).json({ error: `Failed to generate URI: ${error.message}` });
+    } else {
+        res.status(500).json({ error: 'Internal server error during start' });
+    }
   }
 });
 
@@ -80,119 +75,109 @@ app.get('/api/digiid/start', async (req: Request, res: Response) => {
 app.post('/api/digiid/callback', async (req: Request, res: Response) => {
   const { address, uri, signature } = req.body;
 
-  console.log('Received callback:', { address, uri, signature });
-
+  // Basic validation of received data
   if (!address || !uri || !signature) {
-    console.warn('Callback missing required fields.');
-    // Note: DigiID protocol doesn't expect a body response on failure here,
-    // just a non-200 status, but sending JSON for easier debugging.
-    return res.status(400).json({ error: 'Missing required callback parameters.' });
+    console.warn('Callback missing required fields.', { address, uri, signature });
+    // Wallet doesn't expect a body on failure, just non-200. Status only for logging/debug.
+    return res.status(400).send('Missing required callback parameters.');
   }
 
+  const callbackData = { address, uri, signature };
+  console.log('Received callback:', callbackData);
+
+  // --- Nonce Extraction and Session Lookup ---
   let receivedNonce: string | null = null;
   try {
-    // Parse the nonce (parameter 'x') from the received URI
-    const parsedUri = new URL(uri);
+    // DigiID URIs need scheme replaced for standard URL parsing
+    const parsableUri = uri.replace(/^digiid:/, 'http:');
+    const parsedUri = new URL(parsableUri);
     receivedNonce = parsedUri.searchParams.get('x');
   } catch (error) {
     console.warn('Error parsing received URI:', uri, error);
-    return res.status(400).json({ error: 'Invalid URI format.' });
+    return res.status(400).send('Invalid URI format.');
   }
 
   if (!receivedNonce) {
     console.warn('Nonce (x parameter) not found in received URI:', uri);
-    return res.status(400).json({ error: 'Nonce not found in URI.' });
+    return res.status(400).send('Nonce not found in URI.');
   }
 
-  // Find the session ID associated with the nonce
   const sessionId = nonceToSessionMap.get(receivedNonce);
   if (!sessionId) {
-    console.warn('Received nonce does not correspond to any active session:', receivedNonce);
-    // This could happen if the session expired or the nonce is invalid/reused
-    return res.status(404).json({ error: 'Session not found or expired for this nonce.' });
+    console.warn('Session not found for received nonce:', receivedNonce);
+    // Nonce might be expired or invalid
+    return res.status(404).send('Session not found or expired for this nonce.');
   }
 
-  // Retrieve the session state
+  // Retrieve the session *before* the try/finally block for verification
   const session = sessionStore.get(sessionId);
-  if (!session || session.status !== 'pending') {
-    console.warn('Session not found or not in pending state for ID:', sessionId);
-    // Should ideally not happen if nonceToSessionMap is consistent with sessionStore
-    return res.status(404).json({ error: 'Session not found or already completed/failed.' });
+  if (!session) {
+      // This case should be rare if nonceToSessionMap is consistent
+      console.error(`Critical: Session data missing for ${sessionId} despite nonce match.`);
+      return res.status(500).send('Internal server error: Session data missing.');
   }
-
-  // Construct the expected callback URL (must match the one used in /start)
-  const publicUrl = process.env.PUBLIC_URL;
-  if (!publicUrl) {
-     // This should have been caught in /start, but double-check
-     console.error('PUBLIC_URL is unexpectedly missing during callback.');
-     session.status = 'failed';
-     session.error = 'Server configuration error: PUBLIC_URL missing.';
-     return res.status(500).send(); // Send 500 internal server error
+  if (session.status !== 'pending') {
+      console.warn('Session already processed:', sessionId, session.status);
+      // Treat as success here, client will get final status via polling
+      return res.status(200).send('Session already processed.');
   }
+
+
+  // --- Verification ---
   let expectedCallbackUrl: string;
   try {
-     expectedCallbackUrl = new URL('/api/digiid/callback', publicUrl).toString();
+    const publicUrl = process.env.PUBLIC_URL;
+    if (!publicUrl) {
+        // Throw specific error to be caught below
+        throw new Error('PUBLIC_URL environment variable is not configured on the server.');
+    }
+    // Construct expected URL based on server config *at time of verification*
+    expectedCallbackUrl = new URL('/api/digiid/callback', publicUrl).toString();
   } catch (error) {
-     console.error('Invalid PUBLIC_URL format during callback:', publicUrl, error);
-     session.status = 'failed';
-     session.error = 'Server configuration error: Invalid PUBLIC_URL.';
-     return res.status(500).send();
+    // Handle errors during expected URL construction (e.g., invalid PUBLIC_URL)
+    console.error('Server configuration error constructing expected callback URL:', error);
+    session.status = 'failed';
+    session.error = 'Server configuration error preventing verification.';
+    // Update store immediately on this specific failure
+    sessionStore.set(sessionId, session);
+    // Respond 200 OK as per protocol, but status endpoint will show the config error
+    return res.status(200).send();
   }
 
-  const expectedNonce = session.nonce;
+  const verifyOptions = { expectedCallbackUrl, expectedNonce: session.nonce };
 
   try {
-    console.log('Verifying callback with:', {
-      address,
-      uri,
-      signature,
-      expectedCallbackUrl,
-      expectedNonce,
-    });
-
-    // Placeholder for digiidTs.verifyDigiIDCallback call
-    // const isValid = await digiidTs.verifyDigiIDCallback({
-    //   address,
-    //   uri,
-    //   signature,
-    //   callbackUrl: expectedCallbackUrl,
-    //   nonce: expectedNonce,
-    // });
-
-    // --- TEMPORARY PLACEHOLDER VERIFICATION --- 
-    // Simulating verification: check if nonce matches and URI contains expected callback
-    const isValid = receivedNonce === expectedNonce && uri.includes(expectedCallbackUrl);
-    console.log(`Placeholder verification result: ${isValid}`);
-    // --- END PLACEHOLDER --- 
-
-    if (isValid) {
-      console.log(`Verification successful for session ${sessionId}, address: ${address}`);
-      session.status = 'success';
-      session.address = address;
-      // Clean up nonce lookup map once verified successfully
-      nonceToSessionMap.delete(expectedNonce);
-    } else {
-      console.warn(`Verification failed for session ${sessionId}`);
-      session.status = 'failed';
-      session.error = 'Signature verification failed.';
-       // Keep nonce in map for potential debugging, or clean up based on policy
-      // nonceToSessionMap.delete(expectedNonce);
-    }
+    console.log('Attempting to verify callback with:', { callbackData, verifyOptions });
+    // verifyDigiIDCallback throws DigiIDError on failure
+    await verifyDigiIDCallback(callbackData, verifyOptions);
 
-    // Update the session store
-    sessionStore.set(sessionId, session);
-
-    // DigiID protocol expects a 200 OK on success/failure after processing
-    res.status(200).send();
+    // Success case
+    console.log(`Verification successful for session ${sessionId}, address: ${address}`);
+    session.status = 'success';
+    session.address = address; // Store the verified address
+    session.error = undefined; // Clear any previous error
+    nonceToSessionMap.delete(session.nonce); // Clean up nonce map only on success
 
   } catch (error) {
-    console.error('Error during callback verification process for session:', sessionId, error);
-    // Update session state to reflect the error
+    // Failure case (verifyDigiIDCallback threw an error)
+    console.warn(`Verification failed for session ${sessionId}:`, error);
     session.status = 'failed';
-    session.error = 'Internal server error during verification.';
+    if (error instanceof DigiIDError) {
+      session.error = error.message; // Use message from DigiIDError
+    } else if (error instanceof Error) {
+      session.error = `Unexpected verification error: ${error.message}`;
+    } else {
+      session.error = 'An unknown verification error occurred.';
+    }
+    // Optionally cleanup nonce map on failure too, depending on policy
+    // nonceToSessionMap.delete(session.nonce);
+  } finally {
+    // Update store with final status (success/failed) and respond 200 OK
     sessionStore.set(sessionId, session);
-    // Don't expose internal errors via status code if possible, but log them
-    res.status(200).send(); // Still send 200 as processing happened, but log indicates issue
+    console.log(`Final session state for ${sessionId}:`, session);
+    // Wallet expects 200 OK regardless of internal success/fail.
+    // Client uses /status endpoint to get the actual result.
+    res.status(200).send();
   }
 });
 
@@ -200,20 +185,21 @@ app.post('/api/digiid/callback', async (req: Request, res: Response) => {
 app.get('/api/digiid/status/:sessionId', (req: Request, res: Response) => {
   const { sessionId } = req.params;
   const session = sessionStore.get(sessionId);
-
   if (!session) {
+    // Session ID is unknown or expired (and cleaned up)
     return res.status(404).json({ status: 'not_found' });
   }
-
-  // Return only the necessary fields to the client
+  // Return only relevant fields to client
   const { status, address, error } = session;
-  res.json({ status, address, error }); // address and error will be undefined if not set
+  res.json({ status, address, error });
 });
 
+// Simple root endpoint
 app.get('/', (_: Request, res: Response) => {
   res.send('DigiID Demo Backend Running!');
 });
 
+// Start the server
 app.listen(PORT, () => {
   console.log(`Server listening on port ${PORT}`);
-}); 
\ No newline at end of file
+});
\ No newline at end of file