live presence and ui tweaks

Author: jgoux

apps/browser-proxy/src/create-message.ts

@@ -0,0 +1,22 @@
+export function createParameterStatusMessage(name: string, value: string): ArrayBuffer {
+  const encoder = new TextEncoder()
+  const nameBuffer = encoder.encode(name + '\0')
+  const valueBuffer = encoder.encode(value + '\0')
+
+  const messageLength = 4 + nameBuffer.length + valueBuffer.length
+  const message = new ArrayBuffer(1 + messageLength)
+  const view = new DataView(message)
+  const uint8Array = new Uint8Array(message)
+
+  let offset = 0
+  view.setUint8(offset++, 'S'.charCodeAt(0)) // Message type
+  view.setUint32(offset, messageLength, false) // Message length (big-endian)
+  offset += 4
+
+  uint8Array.set(nameBuffer, offset)
+  offset += nameBuffer.length
+
+  uint8Array.set(valueBuffer, offset)
+
+  return message
+}

apps/browser-proxy/src/extract-ip.ts

@@ -0,0 +1,16 @@
+import { isIPv4 } from 'node:net'
+
+export function extractIP(address: string): string {
+  if (isIPv4(address)) {
+    return address
+  }
+
+  // Check if it's an IPv4-mapped IPv6 address
+  const ipv4 = address.match(/::ffff:(\d+\.\d+\.\d+\.\d+)/)
+  if (ipv4) {
+    return ipv4[1]
+  }
+
+  // We assume it's an IPv6 address
+  return address
+}

apps/browser-proxy/src/index.ts

@@ -6,6 +6,8 @@ import makeDebug from 'debug'
 import * as tls from 'node:tls'
 import { extractDatabaseId, isValidServername } from './servername.ts'
 import { getTls } from './tls.ts'
+import { createParameterStatusMessage } from './create-message.ts'
+import { extractIP } from './extract-ip.ts'
 
 const debug = makeDebug('browser-proxy')
 
@@ -127,6 +129,25 @@ tcpServer.on('connection', (socket) => {
     serverVersion() {
       return '16.3'
     },
+    onAuthenticated() {
+      const websocket = websocketConnections.get(databaseId!)
+
+      if (!websocket) {
+        connection.sendError({
+          code: 'XX000',
+          message: 'the browser is not sharing the database',
+          severity: 'FATAL',
+        })
+        connection.end()
+        return
+      }
+
+      const clientIpMessage = createParameterStatusMessage(
+        'client_ip',
+        extractIP(connection.socket.remoteAddress!)
+      )
+      websocket.send(clientIpMessage)
+    },
     onMessage(message, state) {
       if (!state.isAuthenticated) {
         return
@@ -155,6 +176,8 @@ tcpServer.on('connection', (socket) => {
   socket.on('close', () => {
     if (databaseId) {
       tcpConnections.delete(databaseId)
+      const websocket = websocketConnections.get(databaseId)
+      websocket?.send(createParameterStatusMessage('client_ip', ''))
     }
   })
 })

apps/postgres-new/app/db/[id]/page.tsx

@@ -8,7 +8,7 @@ import Workspace from '~/components/workspace'
 export default function Page({ params }: { params: { id: string } }) {
   const databaseId = params.id
   const router = useRouter()
-  const { dbManager, connectedDatabase } = useApp()
+  const { dbManager, liveShare } = useApp()
 
   useEffect(() => {
     async function run() {
@@ -25,12 +25,12 @@ export default function Page({ params }: { params: { id: string } }) {
     run()
   }, [dbManager, databaseId, router])
 
-  // Cleanup connected database when switching databases
+  // Cleanup live shared database when switching databases
   useEffect(() => {
-    if (connectedDatabase.isConnected && connectedDatabase.databaseId !== databaseId) {
-      connectedDatabase.disconnect()
+    if (liveShare.isLiveSharing && liveShare.databaseId !== databaseId) {
+      liveShare.stop()
     }
-  }, [connectedDatabase, databaseId])
+  }, [liveShare, databaseId])
 
   return <Workspace databaseId={databaseId} visibility="local" />
 }

apps/postgres-new/components/app-provider.tsx

@@ -17,6 +17,7 @@ import {
 } from 'react'
 import { DbManager } from '~/lib/db'
 import { useAsyncMemo } from '~/lib/hooks'
+import { parseParameterStatus } from '~/lib/pg-wire-util'
 import { createClient } from '~/utils/supabase/client'
 
 export type AppProps = PropsWithChildren
@@ -105,9 +106,15 @@ export default function AppProvider({ children }: AppProps) {
     return await dbManager.getRuntimePgVersion()
   }, [dbManager])
 
-  const [connectedDatabaseId, setConnectedDatabaseId] = useState<string | null>(null)
-  const [ws, setWs] = useState<WebSocket | null>(null)
-  const connectDatabase = useCallback(
+  const [liveSharedDatabaseId, setLiveSharedDatabaseId] = useState<string | null>(null)
+  const [connectedClientIp, setConnectedClientIp] = useState<string | null>(null)
+  const [liveShareWebsocket, setLiveShareWebsocket] = useState<WebSocket | null>(null)
+  const cleanUp = useCallback(() => {
+    setLiveShareWebsocket(null)
+    setLiveSharedDatabaseId(null)
+    setConnectedClientIp(null)
+  }, [setLiveShareWebsocket, setLiveSharedDatabaseId, setConnectedClientIp])
+  const startLiveShare = useCallback(
     async (databaseId: string) => {
       if (!dbManager) {
         throw new Error('dbManager is not available')
@@ -122,43 +129,53 @@ export default function AppProvider({ children }: AppProps) {
       ws.binaryType = 'arraybuffer'
 
       ws.onopen = () => {
-        setConnectedDatabaseId(databaseId)
+        setLiveSharedDatabaseId(databaseId)
       }
       ws.onmessage = async (event) => {
         const message = new Uint8Array(await event.data)
+
+        const messageType = String.fromCharCode(message[0])
+        if (messageType === 'S') {
+          const { name, value } = parseParameterStatus(message)
+          if (name === 'client_ip') {
+            setConnectedClientIp(value === '' ? null : value)
+            return
+          }
+        }
+
         const response = await db.execProtocolRaw(message)
         ws.send(response)
       }
       ws.onclose = (event) => {
-        setConnectedDatabaseId(null)
+        cleanUp()
       }
       ws.onerror = (error) => {
         console.error('webSocket error:', error)
-        setConnectedDatabaseId(null)
+        cleanUp()
       }
 
-      setWs(ws)
+      setLiveShareWebsocket(ws)
     },
-    [dbManager]
+    [dbManager, cleanUp]
   )
-  const disconnectDatabase = useCallback(() => {
-    ws?.close()
-    setWs(null)
-    setConnectedDatabaseId(null)
-  }, [ws])
-  const connectedDatabase = {
-    connect: connectDatabase,
-    disconnect: disconnectDatabase,
-    databaseId: connectedDatabaseId,
-    isConnected: Boolean(connectedDatabaseId),
+  const stopLiveShare = useCallback(() => {
+    liveShareWebsocket?.close()
+    cleanUp()
+  }, [cleanUp, liveShareWebsocket])
+  const liveShare = {
+    start: startLiveShare,
+    stop: stopLiveShare,
+    databaseId: liveSharedDatabaseId,
+    clientIp: connectedClientIp,
+    isLiveSharing: Boolean(liveSharedDatabaseId),
   }
 
   return (
     <AppContext.Provider
       value={{
         user,
         isLoadingUser,
-        connectedDatabase,
+        liveShare,
         signIn,
         signOut,
         isSignInDialogOpen,
@@ -195,11 +212,12 @@ export type AppContextValues = {
   dbManager?: DbManager
   pgliteVersion?: string
   pgVersion?: string
-  connectedDatabase: {
-    connect: (databaseId: string) => Promise<void>
-    disconnect: () => void
+  liveShare: {
+    start: (databaseId: string) => Promise<void>
+    stop: () => void
     databaseId: string | null
-    isConnected: boolean
+    clientIp: string | null
+    isLiveSharing: boolean
   }
 }
 

apps/postgres-new/components/chat.tsx

@@ -3,7 +3,7 @@
 import { Message, generateId } from 'ai'
 import { useChat } from 'ai/react'
 import { AnimatePresence, m } from 'framer-motion'
-import { ArrowDown, ArrowUp, Flame, Paperclip, Square, WifiOff } from 'lucide-react'
+import { ArrowDown, ArrowUp, Flame, Paperclip, PlugIcon, Square } from 'lucide-react'
 import {
   ChangeEvent,
   FormEventHandler,
@@ -49,7 +49,7 @@ export function getInitialMessages(tables: TablesData): Message[] {
 }
 
 export default function Chat() {
-  const { user, isLoadingUser, focusRef, setIsSignInDialogOpen, isRateLimited, connectedDatabase } =
+  const { user, isLoadingUser, focusRef, setIsSignInDialogOpen, isRateLimited, liveShare } =
     useApp()
   const [inputFocusState, setInputFocusState] = useState(false)
 
@@ -198,7 +198,7 @@ export default function Chat() {
   const [isMessageAnimationComplete, setIsMessageAnimationComplete] = useState(false)
 
   const isChatEnabled =
-    !isLoadingMessages && !isLoadingSchema && user !== undefined && !connectedDatabase.isConnected
+    !isLoadingMessages && !isLoadingSchema && user !== undefined && !liveShare.isLiveSharing
 
   const isSubmitEnabled = isChatEnabled && Boolean(input.trim())
 
@@ -240,31 +240,34 @@ export default function Chat() {
             className={cn(
               'h-full flex flex-col items-center overflow-y-auto',
               !isMessageAnimationComplete ? 'overflow-x-hidden' : undefined,
-              connectedDatabase.isConnected ? 'overflow-y-hidden' : undefined
+              liveShare.isLiveSharing ? 'overflow-y-hidden' : undefined
             )}
             ref={scrollRef}
           >
-            {connectedDatabase.isConnected && (
+            {liveShare.isLiveSharing && (
               <div className="h-full w-full max-w-4xl flex flex-col gap-10 p-10 absolute backdrop-blur-sm bg-card/90">
                 <div className="flex items-center justify-center h-full flex-col gap-y-5">
                   <div className="w-full text-left">
                     <p className="text-lg">Access your in-browser database</p>
                     <p className="text-xs text-muted-foreground">
-                      Closing the window will disconnect your database
+                      Closing the window will stop the Live Share session
                     </p>
                   </div>
                   <CopyableField
-                    value={`postgres://postgres@${connectedDatabase.databaseId}.${process.env.NEXT_PUBLIC_BROWSER_PROXY_DOMAIN}/postgres?sslmode=require`}
+                    value={`postgres://postgres@${liveShare.databaseId}.${process.env.NEXT_PUBLIC_BROWSER_PROXY_DOMAIN}/postgres?sslmode=require`}
                   />
+                  <p className="text-sm text-muted-foreground">
+                    {liveShare.clientIp ? `Connected from ${liveShare.clientIp}` : 'Not connected'}
+                  </p>
                   <Button
                     className="w-full gap-2"
                     variant="outline"
                     onClick={() => {
-                      connectedDatabase.disconnect()
+                      liveShare.stop()
                     }}
                   >
-                    <WifiOff className="h-[1.2rem] w-[1.2rem]" />
-                    <span>Disconnect database</span>
+                    <PlugIcon size={16} />
+                    <span>Stop sharing database</span>
                   </Button>
                 </div>
               </div>

apps/postgres-new/components/live-share-icon.tsx

@@ -0,0 +1,19 @@
+import { cx } from 'class-variance-authority'
+
+export const LiveShareIcon = (props: { size?: number; className?: string }) => (
+  <svg
+    width={props.size}
+    height={props.size}
+    strokeWidth="1"
+    viewBox="0 0 24 24"
+    xmlns="http://www.w3.org/2000/svg"
+    fill="currentColor"
+    className={cx('lucide lucide-live-share', props.className)}
+  >
+    <path
+      fill-rule="evenodd"
+      clip-rule="evenodd"
+      d="M13.735 1.694L15.178 1l8.029 6.328v1.388l-8.029 6.072-1.443-.694v-2.776h-.59c-4.06-.02-6.71.104-10.61 5.163l-1.534-.493a8.23 8.23 0 0 1 .271-2.255 11.026 11.026 0 0 1 3.92-6.793 11.339 11.339 0 0 1 7.502-2.547h1.04v-2.7zm1.804 7.917v2.776l5.676-4.281-5.648-4.545v2.664h-2.86A9.299 9.299 0 0 0 5.77 8.848a10.444 10.444 0 0 0-2.401 4.122c3.351-3.213 6.19-3.359 9.798-3.359h2.373zm-7.647 5.896a4.31 4.31 0 1 1 4.788 7.166 4.31 4.31 0 0 1-4.788-7.166zm.955 5.728a2.588 2.588 0 1 0 2.878-4.302 2.588 2.588 0 0 0-2.878 4.302z"
+    />
+  </svg>
+)