Improve sec

README.md

@@ -29,3 +29,21 @@ pip install -r requirements.txt
 ```bash
 uvicorn app:app.app --reload
 ```
+
+## Credential Retrieval Validation Steps
+
+The `get_credentials` process is designed to verify the identity and authenticity of a validator's request. This involves checking the request's timestamp and verifying its digital signature.
+
+1. **Timestamp Expiry Check**
+   - Ensure the `nonce` (timestamp) in the request is within an acceptable timeframe (`REQUEST_EXPIRY_LIMIT_SECONDS`).
+   - If the request exceeds this time limit, it’s considered expired and is rejected with an error.
+
+2. **Signature Verification**
+   - Construct a message using `validator_info.postfix`, the validator’s unique address, and the `nonce`.
+   - Use the validator's public key to verify that the `signature` provided matches this constructed message.
+   - If the signature is invalid, the request is marked unverified and is rejected.
+
+3. **Error Handling**
+   - If any issues arise—such as an expired request, invalid signature, or missing information—the process returns an error indicating the specific failure reason.
+
+This streamlined validation logic ensures that only timely, authentic requests from verified validators are accepted.

services/image_generation_service.py

@@ -21,6 +21,8 @@
 from fastapi.middleware.cors import CORSMiddleware
 from transformers import AutoTokenizer
 
+REQUEST_EXPIRY_LIMIT_SECONDS = 15
+
 # Define a list of allowed origins (domains)
 allowed_origins = [
     "http://localhost:3000",  # Change this to the domain you want to allow
@@ -33,7 +35,7 @@ def __init__(self, dbhandler, auth_service):
         self.dbhandler = dbhandler
         self.auth_service = auth_service
         self.subtensor = bt.subtensor("finney")
-        self.metagraph = self.subtensor.metagraph(23)
+        self.metagraph: bt.metagraph = self.subtensor.metagraph(23)
 
         self.available_validators = self.dbhandler.get_available_validators()
         self.filter_validators()
@@ -118,9 +120,36 @@ def check_auth(self, key: str) -> None:
     async def get_credentials(
         self, request: Request, validator_info: ValidatorInfo
     ) -> Dict:
+        # Check if validator signed the Request, if not, bypass validating
+        # TODO: remove this logic once all validators updated to latest version
+        if not validator_info.signature:
+            is_verified = True
+            validator_ss58_address = self.metagraph.hotkeys[validator_info.uid]
+        else:
+            # Verify validator by its signature and timestamp
+            try:
+                # Check if the request is within the REQUEST_EXPIRY_LIMIT_SECONDS window
+                received_time_ns = int(validator_info.nonce)
+                current_time_ns = time.time_ns()
+                time_difference_seconds = (current_time_ns - received_time_ns) / 1e9  # Convert nanoseconds to seconds
+
+                if time_difference_seconds > REQUEST_EXPIRY_LIMIT_SECONDS:
+                    raise HTTPException(status_code=400, detail="Request expired")
+
+                # Proceed with signature verification
+                validator_ss58_address = self.metagraph.hotkeys[validator_info.uid]
+                message = f"{validator_info.postfix}{validator_ss58_address}{validator_info.nonce}"
+                keypair = bt.Keypair(ss58_address=validator_ss58_address)
+                is_verified = keypair.verify(message, validator_info.signature)
+            except (ValueError, KeyError, IndexError):
+                is_verified = False
+
+        if not is_verified:
+            raise HTTPException(status_code=400, detail="Cannot verify validator")
+
         client_ip = request.headers.get('X-Real-Ip') or request.client.host
         uid = validator_info.uid
-        hotkey = self.metagraph.hotkeys[uid]
+        hotkey = validator_ss58_address
         postfix = validator_info.postfix
 
         if not postfix:

utils/data_types.py

@@ -40,8 +40,10 @@ class ImageToImage(BaseModel):
 class ValidatorInfo(BaseModel):
     postfix: str
     uid: int
+    # TODO: remove default None once all validators updated to latest version
+    signature: str = None
+    nonce: str = None
     all_uid_info: dict = {}
-    sha: str = ""
 
 class UserSigninInfo(BaseModel):
     email: str