feat: 2FA authorization for web terminal

api/user/auth.go

@@ -86,7 +86,7 @@ func Login(c *gin.Context) {
 	}
 
 	// Check if the user enables 2FA
-	if len(u.OTPSecret) > 0 {
+	if u.EnabledOTP() {
 		if json.OTP == "" && json.RecoveryCode == "" {
 			c.JSON(http.StatusOK, LoginResponse{
 				Message: "The user has enabled 2FA",

api/user/otp.go

@@ -8,6 +8,7 @@ import (
 	"fmt"
 	"github.com/0xJacky/Nginx-UI/api"
 	"github.com/0xJacky/Nginx-UI/internal/crypto"
+	"github.com/0xJacky/Nginx-UI/internal/user"
 	"github.com/0xJacky/Nginx-UI/query"
 	"github.com/0xJacky/Nginx-UI/settings"
 	"github.com/gin-gonic/gin"
@@ -67,8 +68,8 @@ func GenerateTOTP(c *gin.Context) {
 }
 
 func EnrollTOTP(c *gin.Context) {
-	user := api.CurrentUser(c)
-	if len(user.OTPSecret) > 0 {
+	cUser := api.CurrentUser(c)
+	if cUser.EnabledOTP() {
 		c.JSON(http.StatusBadRequest, gin.H{
 			"message": "User already enrolled",
 		})
@@ -109,7 +110,7 @@ func EnrollTOTP(c *gin.Context) {
 	}
 
 	u := query.Auth
-	_, err = u.Where(u.ID.Eq(user.ID)).Update(u.OTPSecret, ciphertext)
+	_, err = u.Where(u.ID.Eq(cUser.ID)).Update(u.OTPSecret, ciphertext)
 	if err != nil {
 		api.ErrHandler(c, err)
 		return
@@ -135,8 +136,8 @@ func ResetOTP(c *gin.Context) {
 		api.ErrHandler(c, err)
 		return
 	}
-	user := api.CurrentUser(c)
-	k := sha1.Sum(user.OTPSecret)
+	cUser := api.CurrentUser(c)
+	k := sha1.Sum(cUser.OTPSecret)
 	if !bytes.Equal(k[:], recoverCode) {
 		c.JSON(http.StatusBadRequest, gin.H{
 			"message": "Invalid recovery code",
@@ -145,7 +146,7 @@ func ResetOTP(c *gin.Context) {
 	}
 
 	u := query.Auth
-	_, err = u.Where(u.ID.Eq(user.ID)).UpdateSimple(u.OTPSecret.Null())
+	_, err = u.Where(u.ID.Eq(cUser.ID)).UpdateSimple(u.OTPSecret.Null())
 	if err != nil {
 		api.ErrHandler(c, err)
 		return
@@ -161,3 +162,40 @@ func OTPStatus(c *gin.Context) {
 		"status": len(api.CurrentUser(c).OTPSecret) > 0,
 	})
 }
+
+func StartSecure2FASession(c *gin.Context) {
+	var json struct {
+		OTP          string `json:"otp"`
+		RecoveryCode string `json:"recovery_code"`
+	}
+	if !api.BindAndValid(c, &json) {
+		return
+	}
+	u := api.CurrentUser(c)
+	if !u.EnabledOTP() {
+		c.JSON(http.StatusBadRequest, gin.H{
+			"message": "User not configured with 2FA",
+		})
+		return
+	}
+
+	if json.OTP == "" && json.RecoveryCode == "" {
+		c.JSON(http.StatusBadRequest, LoginResponse{
+			Message: "The user has enabled 2FA",
+		})
+		return
+	}
+
+	if err := user.VerifyOTP(u, json.OTP, json.RecoveryCode); err != nil {
+		c.JSON(http.StatusBadRequest, LoginResponse{
+			Message: "Invalid 2FA or recovery code",
+		})
+		return
+	}
+
+	sessionId := user.SetSecureSessionID(u.ID)
+
+	c.JSON(http.StatusOK, gin.H{
+		"session_id": sessionId,
+	})
+}

api/user/router.go

@@ -22,5 +22,6 @@ func InitUserRouter(r *gin.RouterGroup) {
 	r.GET("/otp_status", OTPStatus)
 	r.GET("/otp_secret", GenerateTOTP)
 	r.POST("/otp_enroll", EnrollTOTP)
-    r.POST("/otp_reset", ResetOTP)
+	r.POST("/otp_reset", ResetOTP)
+	r.POST("/otp_secure_session", StartSecure2FASession)
 }

app/components.d.ts

@@ -78,8 +78,11 @@ declare module 'vue' {
     NginxControlNginxControl: typeof import('./src/components/NginxControl/NginxControl.vue')['default']
     NodeSelectorNodeSelector: typeof import('./src/components/NodeSelector/NodeSelector.vue')['default']
     NotificationNotification: typeof import('./src/components/Notification/Notification.vue')['default']
+    OTP: typeof import('./src/components/OTP.vue')['default']
     OTPInput: typeof import('./src/components/OTPInput.vue')['default']
     OTPInputOTPInput: typeof import('./src/components/OTPInput/OTPInput.vue')['default']
+    OTPOTPAuthorization: typeof import('./src/components/OTP/OTPAuthorization.vue')['default']
+    OTPOTPAuthorizationModal: typeof import('./src/components/OTP/OTPAuthorizationModal.vue')['default']
     PageHeaderPageHeader: typeof import('./src/components/PageHeader/PageHeader.vue')['default']
     RouterLink: typeof import('vue-router')['RouterLink']
     RouterView: typeof import('vue-router')['RouterView']

app/src/api/otp.ts

@@ -18,6 +18,12 @@ const otp = {
   reset(recovery_code: string) {
     return http.post('/otp_reset', { recovery_code })
   },
+  start_secure_session(passcode: string, recovery_code: string): Promise<{ session_id: string }> {
+    return http.post('/otp_secure_session', {
+      otp: passcode,
+      recovery_code,
+    })
+  },
 }
 
 export default otp

app/src/components/OTP/OTPAuthorization.vue

@@ -0,0 +1,78 @@
+<script setup lang="ts">
+import OTPInput from '@/components/OTPInput/OTPInput.vue'
+
+const emit = defineEmits(['onSubmit'])
+
+const refOTP = ref()
+const useRecoveryCode = ref(false)
+const passcode = ref('')
+const recoveryCode = ref('')
+
+function clickUseRecoveryCode() {
+  passcode.value = ''
+  useRecoveryCode.value = true
+}
+
+function clickUseOTP() {
+  passcode.value = ''
+  useRecoveryCode.value = false
+}
+
+function onSubmit() {
+  emit('onSubmit', passcode.value, recoveryCode.value)
+}
+
+function clearInput() {
+  refOTP.value?.clearInput()
+}
+
+defineExpose({
+  clearInput,
+})
+</script>
+
+<template>
+  <div>
+    <div v-if="!useRecoveryCode">
+      <p>{{ $gettext('Please enter the 2FA code:') }}</p>
+      <OTPInput
+        ref="refOTP"
+        v-model="passcode"
+        class="justify-center mb-6"
+        @on-complete="onSubmit"
+      />
+    </div>
+    <div
+      v-else
+      class="mt-2 mb-4"
+    >
+      <p>{{ $gettext('Input the recovery code:') }}</p>
+      <AInputGroup compact>
+        <AInput v-model:value="recoveryCode" />
+        <AButton
+          type="primary"
+          @click="onSubmit"
+        >
+          {{ $gettext('Recovery') }}
+        </AButton>
+      </AInputGroup>
+    </div>
+
+    <div class="flex justify-center">
+      <a
+        v-if="!useRecoveryCode"
+        @click="clickUseRecoveryCode"
+      >{{ $gettext('Use recovery code') }}</a>
+      <a
+        v-else
+        @click="clickUseOTP"
+      >{{ $gettext('Use OTP') }}</a>
+    </div>
+  </div>
+</template>
+
+<style scoped lang="less">
+:deep(.ant-input-group.ant-input-group-compact) {
+  display: flex;
+}
+</style>

app/src/components/OTP/useOTPModal.ts

@@ -0,0 +1,75 @@
+import { createVNode, render } from 'vue'
+import { Modal, message } from 'ant-design-vue'
+import OTPAuthorization from '@/components/OTP/OTPAuthorization.vue'
+import otp from '@/api/otp'
+
+export interface OTPModalProps {
+  onOk?: (secureSessionId: string) => void
+  onCancel?: () => void
+}
+
+const useOTPModal = () => {
+  const refOTPAuthorization = ref<typeof OTPAuthorization>()
+  const randomId = Math.random().toString(36).substring(2, 8)
+
+  const injectStyles = () => {
+    const style = document.createElement('style')
+
+    style.innerHTML = `
+      .${randomId} .ant-modal-title {
+        font-size: 1.125rem;
+      }
+    `
+    document.head.appendChild(style)
+  }
+
+  const open = ({ onOk, onCancel }: OTPModalProps) => {
+    injectStyles()
+    let container: HTMLDivElement | null = document.createElement('div')
+    document.body.appendChild(container)
+
+    const close = () => {
+      render(null, container!)
+      document.body.removeChild(container!)
+      container = null
+    }
+
+    const verify = (passcode: string, recovery: string) => {
+      otp.start_secure_session(passcode, recovery).then(r => {
+        onOk?.(r.session_id)
+        close()
+      }).catch(async () => {
+        refOTPAuthorization.value?.clearInput()
+        await message.error($gettext('Invalid passcode or recovery code'))
+      })
+    }
+
+    const vnode = createVNode(Modal, {
+      open: true,
+      title: $gettext('Two-factor authentication required'),
+      centered: true,
+      maskClosable: false,
+      class: randomId,
+      footer: false,
+      onCancel: () => {
+        close()
+        onCancel?.()
+      },
+    }, {
+      default: () => h(
+        OTPAuthorization,
+        {
+          ref: refOTPAuthorization,
+          class: 'mt-3',
+          onOnSubmit: verify,
+        },
+      ),
+    })
+
+    render(vnode, container)
+  }
+
+  return { open }
+}
+
+export default useOTPModal

app/src/views/other/Login.vue

@@ -1,14 +1,13 @@
 <script setup lang="ts">
 import { LockOutlined, UserOutlined } from '@ant-design/icons-vue'
 import { Form, message } from 'ant-design-vue'
-import OTPInput from '@/components/OTPInput/OTPInput.vue'
-
 import { useUserStore } from '@/pinia'
 import auth from '@/api/auth'
 import install from '@/api/install'
 import SetLanguage from '@/components/SetLanguage/SetLanguage.vue'
 import SwitchAppearance from '@/components/SwitchAppearance/SwitchAppearance.vue'
 import gettext, { $gettext } from '@/gettext'
+import OTPAuthorization from '@/components/OTP/OTPAuthorization.vue'
 
 const thisYear = new Date().getFullYear()
 
@@ -24,7 +23,6 @@ const loading = ref(false)
 const enabled2FA = ref(false)
 const refOTP = ref()
 const passcode = ref('')
-const useRecoveryCode = ref(false)
 const recoveryCode = ref('')
 
 const modelRef = reactive({
@@ -135,9 +133,13 @@ if (route.query?.code !== undefined && route.query?.state !== undefined) {
   loading.value = false
 }
 
-function clickUseRecoveryCode() {
-  passcode.value = ''
-  useRecoveryCode.value = true
+function handleOTPSubmit(code: string, recovery: string) {
+  passcode.value = code
+  recoveryCode.value = recovery
+
+  nextTick(() => {
+    onSubmit()
+  })
 }
 </script>
 
@@ -173,38 +175,10 @@ function clickUseRecoveryCode() {
               </AFormItem>
             </template>
             <div v-else>
-              <div v-if="!useRecoveryCode">
-                <p>{{ $gettext('Please enter the 2FA code:') }}</p>
-                <OTPInput
-                  ref="refOTP"
-                  v-model="passcode"
-                  class="justify-center mb-6"
-                  @on-complete="onSubmit"
-                />
-
-                <div class="flex justify-center">
-                  <a @click="clickUseRecoveryCode">{{ $gettext('Use recovery code') }}</a>
-                </div>
-              </div>
-
-              <div
-                v-else
-                class="mt-2"
-              >
-                <p>{{ $gettext('Input the recovery code:') }}</p>
-                <AInputGroup compact>
-                  <AInput
-                    v-model:value="recoveryCode"
-                    style="width: calc(100% - 92px)"
-                  />
-                  <AButton
-                    type="primary"
-                    @click="onSubmit"
-                  >
-                    {{ $gettext('Recovery') }}
-                  </AButton>
-                </AInputGroup>
-              </div>
+              <OTPAuthorization
+                ref="refOTP"
+                @on-submit="handleOTPSubmit"
+              />
             </div>
 
             <AFormItem v-if="!enabled2FA">

app/src/views/pty/Terminal.vue

@@ -2,20 +2,43 @@
 import '@xterm/xterm/css/xterm.css'
 import { Terminal } from '@xterm/xterm'
 import { FitAddon } from '@xterm/addon-fit'
-import { onMounted, onUnmounted } from 'vue'
 import _ from 'lodash'
 import ws from '@/lib/websocket'
+import useOTPModal from '@/components/OTP/useOTPModal'
 
 let term: Terminal | null
 let ping: NodeJS.Timeout
 
-const websocket = ws('/api/pty')
+const router = useRouter()
+const websocket = shallowRef()
+const lostConnection = ref(false)
 
 onMounted(() => {
-  initTerm()
-
-  websocket.onmessage = wsOnMessage
-  websocket.onopen = wsOnOpen
+  const otpModal = useOTPModal()
+
+  otpModal.open({
+    onOk(secureSessionId: string) {
+      websocket.value = ws(`/api/pty?X-Secure-Session-ID=${secureSessionId}`, false)
+
+      nextTick(() => {
+        initTerm()
+        websocket.value.onmessage = wsOnMessage
+        websocket.value.onopen = wsOnOpen
+        websocket.value.onerror = () => {
+          lostConnection.value = true
+        }
+        websocket.value.onclose = () => {
+          lostConnection.value = true
+        }
+      })
+    },
+    onCancel() {
+      if (window.history.length > 1)
+        router.go(-1)
+      else
+        router.push('/')
+    },
+  })
 })
 
 interface Message {
@@ -65,7 +88,7 @@ function initTerm() {
 }
 
 function sendMessage(data: Message) {
-  websocket.send(JSON.stringify(data))
+  websocket.value.send(JSON.stringify(data))
 }
 
 function wsOnMessage(msg: { data: string | Uint8Array }) {
@@ -82,13 +105,20 @@ onUnmounted(() => {
   window.removeEventListener('resize', fit)
   clearInterval(ping)
   term?.dispose()
-  websocket.close()
+  websocket.value?.close()
 })
 
 </script>
 
 <template>
   <ACard :title="$gettext('Terminal')">
+    <AAlert
+      v-if="lostConnection"
+      class="mb-6"
+      type="error"
+      show-icon
+      :message="$gettext('Connection lost, please refresh the page.')"
+    />
     <div
       id="terminal"
       class="console"

internal/cache/cache.go

@@ -0,0 +1,31 @@
+package cache
+
+import (
+	"github.com/0xJacky/Nginx-UI/internal/logger"
+	"github.com/dgraph-io/ristretto"
+	"time"
+)
+
+var cache *ristretto.Cache
+
+func Init() {
+	var err error
+	cache, err = ristretto.NewCache(&ristretto.Config{
+		NumCounters: 1e7,     // number of keys to track frequency of (10M).
+		MaxCost:     1 << 30, // maximum cost of cache (1GB).
+		BufferItems: 64,      // number of keys per Get buffer.
+	})
+
+	if err != nil {
+		logger.Fatal("initializing local cache err", err)
+	}
+}
+
+func Set(key interface{}, value interface{}, ttl time.Duration) {
+	cache.SetWithTTL(key, value, 0, ttl)
+	cache.Wait()
+}
+
+func Get(key interface{}) (value interface{}, ok bool) {
+	return cache.Get(key)
+}

internal/kernal/boot.go

@@ -4,6 +4,7 @@ import (
 	"crypto/rand"
 	"encoding/hex"
 	"github.com/0xJacky/Nginx-UI/internal/analytic"
+	"github.com/0xJacky/Nginx-UI/internal/cache"
 	"github.com/0xJacky/Nginx-UI/internal/cert"
 	"github.com/0xJacky/Nginx-UI/internal/cluster"
 	"github.com/0xJacky/Nginx-UI/internal/cron"
@@ -26,6 +27,7 @@ func Boot() {
 		InitNodeSecret,
 		InitCryptoSecret,
 		validation.Init,
+		cache.Init,
 	}
 
 	syncs := []func(){

internal/user/otp.go

@@ -4,10 +4,14 @@ import (
 	"bytes"
 	"crypto/sha1"
 	"encoding/hex"
+	"fmt"
+	"github.com/0xJacky/Nginx-UI/internal/cache"
 	"github.com/0xJacky/Nginx-UI/internal/crypto"
 	"github.com/0xJacky/Nginx-UI/model"
+	"github.com/google/uuid"
 	"github.com/pkg/errors"
 	"github.com/pquerna/otp/totp"
+	"time"
 )
 
 var (
@@ -37,3 +41,23 @@ func VerifyOTP(user *model.Auth, otp, recoveryCode string) (err error) {
 	}
 	return
 }
+
+func secureSessionIDCacheKey(sessionId string) string {
+	return fmt.Sprintf("otp_secure_session:_%s", sessionId)
+}
+
+func SetSecureSessionID(userId int) (sessionId string) {
+	sessionId = uuid.NewString()
+	cache.Set(secureSessionIDCacheKey(sessionId), userId, 5*time.Minute)
+
+	return
+}
+
+func VerifySecureSessionID(sessionId string, userId int) bool {
+	if v, ok := cache.Get(secureSessionIDCacheKey(sessionId)); ok {
+		if v.(int) == userId {
+			return true
+		}
+	}
+	return false
+}

model/auth.go

@@ -14,3 +14,7 @@ type AuthToken struct {
 	Token     string `json:"token"`
 	ExpiredAt int64  `json:"expired_at" gorm:"default:0"`
 }
+
+func (u *Auth) EnabledOTP() bool {
+	return len(u.OTPSecret) != 0
+}

router/middleware.go

@@ -5,6 +5,7 @@ import (
 	"github.com/0xJacky/Nginx-UI/app"
 	"github.com/0xJacky/Nginx-UI/internal/logger"
 	"github.com/0xJacky/Nginx-UI/internal/user"
+	"github.com/0xJacky/Nginx-UI/model"
 	"github.com/0xJacky/Nginx-UI/settings"
 	"github.com/gin-contrib/static"
 	"github.com/gin-gonic/gin"
@@ -74,6 +75,41 @@ func authRequired() gin.HandlerFunc {
 	}
 }
 
+func required2FA() gin.HandlerFunc {
+	return func(c *gin.Context) {
+		u, ok := c.Get("user")
+		if !ok {
+			c.Next()
+			return
+		}
+		cUser := u.(*model.Auth)
+		if !cUser.EnabledOTP() {
+			c.Next()
+			return
+		}
+		ssid := c.GetHeader("X-Secure-Session-ID")
+		if ssid == "" {
+			ssid = c.Query("X-Secure-Session-ID")
+		}
+		if ssid == "" {
+			c.AbortWithStatusJSON(http.StatusForbidden, gin.H{
+				"message": "Secure Session ID is empty",
+			})
+			return
+		}
+
+		if user.VerifySecureSessionID(ssid, cUser.ID) {
+			c.Next()
+			return
+		}
+
+		c.AbortWithStatusJSON(http.StatusForbidden, gin.H{
+			"message": "Secure Session ID is invalid",
+		})
+		return
+	}
+}
+
 type serverFileSystemType struct {
 	http.FileSystem
 }

router/routers.go

@@ -69,7 +69,10 @@ func InitRouter() *gin.Engine {
 		{
 			analytic.InitWebSocketRouter(w)
 			certificate.InitCertificateWebSocketRouter(w)
-			terminal.InitRouter(w)
+			o := w.Group("", required2FA())
+			{
+				terminal.InitRouter(o)
+			}
 			nginx.InitNginxLogRouter(w)
 			upstream.InitRouter(w)
 			system.InitWebSocketRouter(w)