Improve reading job tokens from the environment

inc/TJAlienConnectionManager.h

@@ -68,7 +68,7 @@ private:
     static std::string readBuffer;
 
 public:
-    TJAlienConnectionManager();
+    TJAlienConnectionManager() {} // default constructor
     ~TJAlienConnectionManager();
     int CreateConnection();
     void ConnectJBox(TJAlienCredentialsObject c);

inc/TJAlienCredentials.h

@@ -16,46 +16,39 @@ enum CredentialsKind { cJBOX_TOKEN = 0,
 };
 
 
-struct TJAlienCredentialsObject {
+class TJAlienCredentialsObject {
+ public:
     string certpath;
     string keypath;
-    string source;
     string password;
     CredentialsKind kind;
+    bool autoremove;
 
     TJAlienCredentialsObject() {}
     TJAlienCredentialsObject(string certpath,
                              string keypath,
                              CredentialsKind kind = cOTHER_TOKEN,
-                             string source = "")
+                             bool autoremove = false)
     {
         this->certpath = certpath;
         this->keypath = keypath;
         this->kind = kind;
-        this->source = source;
+        this->autoremove = autoremove;
     };
 
-    bool exists() { return fileExists(certpath) && fileExists(keypath); }
+    void wipe() {
+        if(autoremove) {
+            if(gDebug) printf("removing safe files: %s %s\n", certpath.c_str(), keypath.c_str());
+            remove(certpath.c_str());
+            remove(keypath.c_str());
+        }
+    }
+
+    bool exists();
     const string getKey();
     const string getCertificate();
     const string getPassword();
     void readPassword();
-
-private:
-    bool fileExists(string filename)
-    {
-        bool fileExists = false;
-        FILE *f = fopen(filename.c_str(), "r");
-
-        if (f != NULL ) {
-            fclose(f);
-            fileExists = true;
-        } else {
-            fileExists = false;
-        }
-
-        return fileExists;
-    }
 };
 
 class TJAlienCredentials : public TObject {
@@ -63,8 +56,8 @@ public:
   TJAlienCredentials();
   ~TJAlienCredentials();
 
-  string getTmpDir();
-  string getHomeDir();
+  static string getTmpDir();
+  static string getHomeDir();
   void loadCredentials();
 
   bool has(CredentialsKind kind);
@@ -74,8 +67,8 @@ public:
 
   static const char *ENV_JOBTOKEN_KEY;
   static const char *ENV_JOBTOKEN_CERT;
-  static const char *TMP_JOBTOKEN_KEY_FNAME;
-  static const char *TMP_JOBTOKEN_CERT_FNAME;
+  static const char *TMP_JOBTOKEN_KEY_FNAME_PREFIX;
+  static const char *TMP_JOBTOKEN_CERT_FNAME_PREFIX;
 
 private:
   void loadTokenCertificate();
@@ -86,6 +79,9 @@ private:
   string getTokencertPath();
   string getTokenkeyPath();
 
+  string getSafeFilename(const string& prefix);
+  void writeSafeFile(const string& filepath, const string& content);
+
   string tmpdir;
   string homedir;
   map<CredentialsKind, TJAlienCredentialsObject> found_credentials;

src/TJAlienConnectionManager.cxx

@@ -12,13 +12,13 @@ int TJAlienConnectionManager::writeable_flag = 0;
 int TJAlienConnectionManager::receive_flag = 0;
 std::string TJAlienConnectionManager::readBuffer = "";
 
-TJAlienConnectionManager::TJAlienConnectionManager() {
-    creds.loadCredentials();
-}
-
 TJAlienConnectionManager::~TJAlienConnectionManager() {
     if (context)
         lws_context_destroy(context);
+
+    if (creds.has(cJOB_TOKEN)) {
+        creds.removeCredentials(cJOB_TOKEN);
+    }
 }
 
 //______________________________________________________________________________
@@ -30,6 +30,7 @@ int TJAlienConnectionManager::CreateConnection()
 
     clearFlags();
 
+    creds.loadCredentials();
     if (creds.count() == 0) {
         Error("TJAlienConnectionManager", "Failed to get any credentials");
         return -1;
@@ -247,6 +248,10 @@ void TJAlienConnectionManager::MakeWebsocketConnection(TJAlienCredentialsObject
         }
     }
 
+    if(creds.kind == cJOB_TOKEN) {
+        this->creds.removeCredentials(creds.kind);
+    }
+
     creation_info.ssl_private_key_password = "";
     fWSPort = WSPort;
     return;

src/TJAlienCredentials.cxx

@@ -5,6 +5,7 @@
 #include <fstream>
 #include <cstdlib>
 #include <iostream>
+#include <fcntl.h>
 
 using std::ifstream;
 using std::ofstream;
@@ -14,8 +15,56 @@ using std::getenv;
 
 const char* TJAlienCredentials::ENV_JOBTOKEN_KEY = "JALIEN_TOKEN_KEY";
 const char* TJAlienCredentials::ENV_JOBTOKEN_CERT = "JALIEN_TOKEN_CERT";
-const char* TJAlienCredentials::TMP_JOBTOKEN_KEY_FNAME = "tmpjobtokenkey.pem";
-const char* TJAlienCredentials::TMP_JOBTOKEN_CERT_FNAME = "tmpjobtokencert.pem";
+
+const char* TJAlienCredentials::TMP_JOBTOKEN_KEY_FNAME_PREFIX = "tmpjobtokenkey_";
+const char* TJAlienCredentials::TMP_JOBTOKEN_CERT_FNAME_PREFIX = "tmpjobtokencert_";
+
+bool fileExists(const string &filename)
+{
+    bool fileExists = false;
+    FILE *f = fopen(filename.c_str(), "r");
+
+    if (f != NULL ) {
+        fclose(f);
+        fileExists = true;
+    } else {
+        fileExists = false;
+    }
+
+    return fileExists;
+}
+
+
+bool TJAlienCredentialsObject::exists() {
+    return fileExists(certpath) && fileExists(keypath);
+}
+
+void TJAlienCredentials::writeSafeFile(const string& filename, const string& content) {
+    if(gDebug) printf("writing safe file %s\n", filename.c_str());
+    int fd = open(filename.c_str(), O_RDWR | O_CREAT, 0600);
+    write(fd, content.c_str(), content.length());
+    close(fd);
+}
+
+string TJAlienCredentials::getSafeFilename(const string& prefix) {
+    string filename = TJAlienCredentials::getTmpDir() + "/" + prefix;
+
+    const char *JOB_ID = getenv("ALIEN_PROC_ID");
+
+    if(JOB_ID != NULL) {
+        filename += JOB_ID;
+    } else {
+        pid_t pid = getpid();
+        unsigned int rnd = random() % 100 + 1;
+
+        do {
+            filename += std::to_string(pid) + "_" + std::to_string(rnd);
+        } while(fileExists(filename));
+    }
+
+
+    return filename;
+}
 
 string TJAlienCredentials::getTmpDir() {
 	string tmpdir;
@@ -83,6 +132,7 @@ TJAlienCredentials::TJAlienCredentials() {
 }
 
 void TJAlienCredentials::loadCredentials() {
+  removeCredentials(cJOB_TOKEN);
 	found_credentials.clear();
 	loadTokenCertificate();
 	loadFullGridCertificate();
@@ -114,24 +164,18 @@ void TJAlienCredentials::loadJobTokenCertificate() {
         return;
     }
 
-    string tmpcertpath = getTmpDir() + "/" + TMP_JOBTOKEN_CERT_FNAME;
-    ofstream certFile(tmpcertpath);
-    certFile << env_cert;
-    certFile.close();
+    // environment variables contain valid filepaths instead of the actual token
+    if(fileExists(env_cert) && fileExists(env_key)) {
+        found_credentials[cJOB_TOKEN] = TJAlienCredentialsObject(env_cert, env_key, cJOB_TOKEN);
+    } else {
+        const string& tmpcertpath = getSafeFilename(TMP_JOBTOKEN_CERT_FNAME_PREFIX);
+        writeSafeFile(tmpcertpath, env_cert);
 
-    string tmpkeypath = getTmpDir() + "/" + TMP_JOBTOKEN_KEY_FNAME;
-    ofstream keyFile(tmpkeypath);
-    keyFile << env_key;
-    keyFile.close();
+        const string& tmpkeypath = getSafeFilename(TMP_JOBTOKEN_KEY_FNAME_PREFIX);
+        writeSafeFile(tmpkeypath, env_key);
 
-    TJAlienCredentialsObject jobtoken_certificate(tmpcertpath, tmpkeypath, cJOB_TOKEN);
-    if(jobtoken_certificate.exists()) {
-        found_credentials[cJOB_TOKEN] = jobtoken_certificate;
-    } else {
-        remove(tmpkeypath.c_str());
-        remove(tmpcertpath.c_str());
+        found_credentials[cJOB_TOKEN] = TJAlienCredentialsObject(tmpcertpath, tmpkeypath, cJOB_TOKEN, true);
     }
-
 }
 
 bool TJAlienCredentials::has(CredentialsKind kind) {
@@ -148,6 +192,9 @@ TJAlienCredentialsObject TJAlienCredentials::get(CredentialsKind kind) {
 
 void TJAlienCredentials::removeCredentials(CredentialsKind kind) {
 	if (this->has(kind)) {
+      if(kind == cJOB_TOKEN)
+          get(kind).wipe();
+
 		found_credentials.erase(kind);
 	}
 }
@@ -202,14 +249,9 @@ const string TJAlienCredentialsObject::getPassword() {
 	if (this->password.empty())
 		readPassword();
 
-	//printf("this->password %s\n", this->password.c_str());
     return this->password;
 }
 
 TJAlienCredentials::~TJAlienCredentials() {
-    if(has(cJOB_TOKEN)) {
-        TJAlienCredentialsObject creds = get(cJOB_TOKEN);
-        remove(creds.certpath.c_str());
-        remove(creds.keypath.c_str());
-    }
+    removeCredentials(cJOB_TOKEN);
 }