]> git.netwichtig.de Git - user/henk/code/inspircd.git/blob - src/modules/extra/m_pgsql.cpp
5d07c23b5a9b10a167bb1039ca1851b4c64435b3
[user/henk/code/inspircd.git] / src / modules / extra / m_pgsql.cpp
1 /*       +------------------------------------+
2  *       | Inspire Internet Relay Chat Daemon |
3  *       +------------------------------------+
4  *
5  *  InspIRCd is copyright (C) 2002-2004 ChatSpike-Dev.
6  *                       E-mail:
7  *                <brain@chatspike.net>
8  *                <Craig@chatspike.net>
9  *                <omster@gmail.com>
10  *     
11  * Written by Craig Edwards, Craig McLure, and others.
12  * This program is free but copyrighted software; see
13  *            the file COPYING for details.
14  *
15  * ---------------------------------------------------
16  */
17
18 #include <cstdlib>
19 #include <sstream>
20 #include <string>
21 #include <deque>
22 #include <map>
23 #include <libpq-fe.h>
24
25 #include "users.h"
26 #include "channels.h"
27 #include "modules.h"
28 #include "helperfuncs.h"
29 #include "inspircd.h"
30 #include "configreader.h"
31
32 #include "m_sqlv2.h"
33
34 /* $ModDesc: PostgreSQL Service Provider module for all other m_sql* modules, uses v2 of the SQL API */
35 /* $CompileFlags: -I`pg_config --includedir` `perl extra/pgsql_config.pl` */
36 /* $LinkerFlags: -L`pg_config --libdir` -lpq */
37
38 /* UGH, UGH, UGH, UGH, UGH, UGH
39  * I'm having trouble seeing how I
40  * can avoid this. The core-defined
41  * constructors for InspSocket just
42  * aren't suitable...and if I'm
43  * reimplementing them I need this so
44  * I can access the socket engine :\
45  */
46 extern InspIRCd* ServerInstance;
47 extern InspSocket* socket_ref[MAX_DESCRIPTORS];
48 extern time_t TIME;
49
50 /* Forward declare, so we can have the typedef neatly at the top */
51 class SQLConn;
52 /* Also needs forward declaration, as it's used inside SQLconn */
53 class ModulePgSQL;
54
55 typedef std::map<std::string, SQLConn*> ConnMap;
56
57 /* CREAD,       Connecting and wants read event
58  * CWRITE,      Connecting and wants write event
59  * WREAD,       Connected/Working and wants read event
60  * WWRITE,      Connected/Working and wants write event
61  * RREAD,       Resetting and wants read event
62  * RWRITE,      Resetting and wants write event
63  */
64 enum SQLstatus { CREAD, CWRITE, WREAD, WWRITE, RREAD, RWRITE };
65
66 /** SQLhost, simple structure to store information about a SQL-connection-to-be
67  * We use this struct simply to make it neater passing around host information
68  * when we're creating connections and resolving hosts.
69  * Rather than giving SQLresolver a parameter for every field here so it can in
70  * turn call SQLConn's constructor with them all, both can simply use a SQLhost.
71  */
72 class SQLhost
73 {
74  public:
75         std::string             id;             /* Database handle id */
76         std::string     host;   /* Database server hostname */
77         unsigned int    port;   /* Database server port */
78         std::string     name;   /* Database name */
79         std::string     user;   /* Database username */
80         std::string     pass;   /* Database password */
81         bool                    ssl;    /* If we should require SSL */
82  
83         SQLhost()
84         {
85         }               
86
87         SQLhost(const std::string& i, const std::string& h, unsigned int p, const std::string& n, const std::string& u, const std::string& pa, bool s)
88         : id(i), host(h), port(p), name(n), user(u), pass(pa), ssl(s)
89         {
90         }
91 };
92
93 class SQLresolver : public Resolver
94 {
95  private:
96         SQLhost host;
97         ModulePgSQL* mod;
98  public:
99         SQLresolver(ModulePgSQL* m, Server* srv, const SQLhost& hi)
100         : Resolver(hi.host, DNS_QUERY_FORWARD), host(hi), mod(m)
101         {
102         }
103
104         virtual void OnLookupComplete(const std::string &result);
105
106         virtual void OnError(ResolverError e, const std::string &errormessage)
107         {
108                 log(DEBUG, "DNS lookup failed (%s), dying horribly", errormessage.c_str());
109         }
110
111         virtual ~SQLresolver()
112         {
113         }
114 };
115
116 /** QueryQueue, a queue of queries waiting to be executed.
117  * This maintains two queues internally, one for 'priority'
118  * queries and one for less important ones. Each queue has
119  * new queries appended to it and ones to execute are popped
120  * off the front. This keeps them flowing round nicely and no
121  * query should ever get 'stuck' for too long. If there are
122  * queries in the priority queue they will be executed first,
123  * 'unimportant' queries will only be executed when the
124  * priority queue is empty.
125  *
126  * We store lists of SQLrequest's here, by value as we want to avoid storing
127  * any data allocated inside the client module (in case that module is unloaded
128  * while the query is in progress).
129  *
130  * Because we want to work on the current SQLrequest in-situ, we need a way
131  * of accessing the request we are currently processing, QueryQueue::front(),
132  * but that call needs to always return the same request until that request
133  * is removed from the queue, this is what the 'which' variable is. New queries are
134  * always added to the back of one of the two queues, but if when front()
135  * is first called then the priority queue is empty then front() will return
136  * a query from the normal queue, but if a query is then added to the priority
137  * queue then front() must continue to return the front of the *normal* queue
138  * until pop() is called.
139  */
140
141 class QueryQueue : public classbase
142 {
143 private:
144         typedef std::deque<SQLrequest> ReqDeque;        
145
146         ReqDeque priority;      /* The priority queue */
147         ReqDeque normal;        /* The 'normal' queue */
148         enum { PRI, NOR, NON } which;   /* Which queue the currently active element is at the front of */
149
150 public:
151         QueryQueue()
152         : which(NON)
153         {
154         }
155         
156         void push(const SQLrequest &q)
157         {
158                 log(DEBUG, "QueryQueue::push(): Adding %s query to queue: %s", ((q.pri) ? "priority" : "non-priority"), q.query.q.c_str());
159                 
160                 if(q.pri)
161                         priority.push_back(q);
162                 else
163                         normal.push_back(q);
164         }
165         
166         void pop()
167         {
168                 if((which == PRI) && priority.size())
169                 {
170                         priority.pop_front();
171                 }
172                 else if((which == NOR) && normal.size())
173                 {
174                         normal.pop_front();
175                 }
176                 
177                 /* Reset this */
178                 which = NON;
179                 
180                 /* Silently do nothing if there was no element to pop() */
181         }
182         
183         SQLrequest& front()
184         {
185                 switch(which)
186                 {
187                         case PRI:
188                                 return priority.front();
189                         case NOR:
190                                 return normal.front();
191                         default:
192                                 if(priority.size())
193                                 {
194                                         which = PRI;
195                                         return priority.front();
196                                 }
197                                 
198                                 if(normal.size())
199                                 {
200                                         which = NOR;
201                                         return normal.front();
202                                 }
203                                 
204                                 /* This will probably result in a segfault,
205                                  * but the caller should have checked totalsize()
206                                  * first so..meh - moron :p
207                                  */
208                                 
209                                 return priority.front();
210                 }
211         }
212         
213         std::pair<int, int> size()
214         {
215                 return std::make_pair(priority.size(), normal.size());
216         }
217         
218         int totalsize()
219         {
220                 return priority.size() + normal.size();
221         }
222         
223         void PurgeModule(Module* mod)
224         {
225                 DoPurgeModule(mod, priority);
226                 DoPurgeModule(mod, normal);
227         }
228         
229 private:
230         void DoPurgeModule(Module* mod, ReqDeque& q)
231         {
232                 for(ReqDeque::iterator iter = q.begin(); iter != q.end(); iter++)
233                 {
234                         if(iter->GetSource() == mod)
235                         {
236                                 if(iter->id == front().id)
237                                 {
238                                         /* It's the currently active query.. :x */
239                                         iter->SetSource(NULL);
240                                 }
241                                 else
242                                 {
243                                         /* It hasn't been executed yet..just remove it */
244                                         iter = q.erase(iter);
245                                 }
246                         }
247                 }
248         }
249 };
250
251 /** PgSQLresult is a subclass of the mostly-pure-virtual class SQLresult.
252  * All SQL providers must create their own subclass and define it's methods using that
253  * database library's data retriveal functions. The aim is to avoid a slow and inefficient process
254  * of converting all data to a common format before it reaches the result structure. This way
255  * data is passes to the module nearly as directly as if it was using the API directly itself.
256  */
257
258 class PgSQLresult : public SQLresult
259 {
260         PGresult* res;
261         int currentrow;
262         int rows;
263         int cols;
264         
265         SQLfieldList* fieldlist;
266         SQLfieldMap* fieldmap;
267 public:
268         PgSQLresult(Module* self, Module* to, unsigned long id, PGresult* result)
269         : SQLresult(self, to, id), res(result), currentrow(0), fieldlist(NULL), fieldmap(NULL)
270         {
271                 rows = PQntuples(res);
272                 cols = PQnfields(res);
273                 
274                 log(DEBUG, "Created new PgSQL result; %d rows, %d columns, %s affected", rows, cols, PQcmdTuples(res));
275         }
276         
277         ~PgSQLresult()
278         {
279                 /* If we allocated these, free them... */
280                 if(fieldlist)
281                         DELETE(fieldlist);
282                 
283                 if(fieldmap)
284                         DELETE(fieldmap);
285                 
286                 PQclear(res);
287         }
288         
289         virtual int Rows()
290         {
291                 if(!cols && !rows)
292                 {
293                         return atoi(PQcmdTuples(res));
294                 }
295                 else
296                 {
297                         return rows;
298                 }
299         }
300         
301         virtual int Cols()
302         {
303                 return PQnfields(res);
304         }
305         
306         virtual std::string ColName(int column)
307         {
308                 char* name = PQfname(res, column);
309                 
310                 return (name) ? name : "";
311         }
312         
313         virtual int ColNum(const std::string &column)
314         {
315                 int n = PQfnumber(res, column.c_str());
316                 
317                 if(n == -1)
318                 {
319                         throw SQLbadColName();
320                 }
321                 else
322                 {
323                         return n;
324                 }
325         }
326         
327         virtual SQLfield GetValue(int row, int column)
328         {
329                 char* v = PQgetvalue(res, row, column);
330                 
331                 if(v)
332                 {
333                         return SQLfield(std::string(v, PQgetlength(res, row, column)), PQgetisnull(res, row, column));
334                 }
335                 else
336                 {
337                         log(DEBUG, "PQgetvalue returned a null pointer..nobody wants to tell us what this means");
338                         throw SQLbadColName();
339                 }
340         }
341         
342         virtual SQLfieldList& GetRow()
343         {
344                 /* In an effort to reduce overhead we don't actually allocate the list
345                  * until the first time it's needed...so...
346                  */
347                 if(fieldlist)
348                 {
349                         fieldlist->clear();
350                 }
351                 else
352                 {
353                         fieldlist = new SQLfieldList;
354                 }
355                 
356                 if(currentrow < PQntuples(res))
357                 {
358                         int cols = PQnfields(res);
359                         
360                         for(int i = 0; i < cols; i++)
361                         {
362                                 fieldlist->push_back(GetValue(currentrow, i));
363                         }
364                         
365                         currentrow++;
366                 }
367                 
368                 return *fieldlist;
369         }
370         
371         virtual SQLfieldMap& GetRowMap()
372         {
373                 /* In an effort to reduce overhead we don't actually allocate the map
374                  * until the first time it's needed...so...
375                  */
376                 if(fieldmap)
377                 {
378                         fieldmap->clear();
379                 }
380                 else
381                 {
382                         fieldmap = new SQLfieldMap;
383                 }
384                 
385                 if(currentrow < PQntuples(res))
386                 {
387                         int cols = PQnfields(res);
388                         
389                         for(int i = 0; i < cols; i++)
390                         {
391                                 fieldmap->insert(std::make_pair(ColName(i), GetValue(currentrow, i)));
392                         }
393                         
394                         currentrow++;
395                 }
396                 
397                 return *fieldmap;
398         }
399         
400         virtual SQLfieldList* GetRowPtr()
401         {
402                 SQLfieldList* fl = new SQLfieldList;
403                 
404                 if(currentrow < PQntuples(res))
405                 {
406                         int cols = PQnfields(res);
407                         
408                         for(int i = 0; i < cols; i++)
409                         {
410                                 fl->push_back(GetValue(currentrow, i));
411                         }
412                         
413                         currentrow++;
414                 }
415                 
416                 return fl;
417         }
418         
419         virtual SQLfieldMap* GetRowMapPtr()
420         {
421                 SQLfieldMap* fm = new SQLfieldMap;
422                 
423                 if(currentrow < PQntuples(res))
424                 {
425                         int cols = PQnfields(res);
426                         
427                         for(int i = 0; i < cols; i++)
428                         {
429                                 fm->insert(std::make_pair(ColName(i), GetValue(currentrow, i)));
430                         }
431                         
432                         currentrow++;
433                 }
434                 
435                 return fm;
436         }
437         
438         virtual void Free(SQLfieldMap* fm)
439         {
440                 DELETE(fm);
441         }
442         
443         virtual void Free(SQLfieldList* fl)
444         {
445                 DELETE(fl);
446         }
447 };
448
449 /** SQLConn represents one SQL session.
450  * Each session has its own persistent connection to the database.
451  * This is a subclass of InspSocket so it can easily recieve read/write events from the core socket
452  * engine, unlike the original MySQL module this module does not block. Ever. It gets a mild stabbing
453  * if it dares to.
454  */
455
456 class SQLConn : public InspSocket
457 {
458 private:
459         ModulePgSQL* us;                /* Pointer to the SQL provider itself */
460         Server* Srv;                    /* Server* for..uhm..something, maybe */
461         std::string     dbhost; /* Database server hostname */
462         unsigned int    dbport; /* Database server port */
463         std::string     dbname; /* Database name */
464         std::string     dbuser; /* Database username */
465         std::string     dbpass; /* Database password */
466         bool                    ssl;    /* If we should require SSL */
467         PGconn*                 sql;    /* PgSQL database connection handle */
468         SQLstatus               status; /* PgSQL database connection status */
469         bool                    qinprog;/* If there is currently a query in progress */
470         QueryQueue              queue;  /* Queue of queries waiting to be executed on this connection */
471         time_t                  idle;   /* Time we last heard from the database */
472
473 public:
474
475         /* This class should only ever be created inside this module, using this constructor, so we don't have to worry about the default ones */
476
477         SQLConn(ModulePgSQL* self, Server* srv, const SQLhost& hostinfo);
478
479         ~SQLConn();
480
481         bool DoConnect();
482
483         virtual void Close();
484         
485         bool DoPoll();
486         
487         bool DoConnectedPoll();
488
489         bool DoResetPoll();
490         
491         void ShowStatus();      
492         
493         virtual bool OnDataReady();
494
495         virtual bool OnWriteReady();
496         
497         virtual bool OnConnected();
498         
499         bool DoEvent();
500         
501         bool Reconnect();
502         
503         std::string MkInfoStr();
504         
505         const char* StatusStr();
506         
507         SQLerror DoQuery(SQLrequest &req);
508         
509         SQLerror Query(const SQLrequest &req);
510         
511         void OnUnloadModule(Module* mod);
512 };
513
514 class ModulePgSQL : public Module
515 {
516 private:
517         Server* Srv;
518         ConnMap connections;
519         unsigned long currid;
520         char* sqlsuccess;
521
522 public:
523         ModulePgSQL(Server* Me)
524         : Module::Module(Me), Srv(Me), currid(0)
525         {
526                 log(DEBUG, "%s 'SQL' feature", Srv->PublishFeature("SQL", this) ? "Published" : "Couldn't publish");
527                 
528                 sqlsuccess = new char[strlen(SQLSUCCESS)+1];
529                 
530                 strcpy(sqlsuccess, SQLSUCCESS);
531
532                 OnRehash("");
533         }
534
535         void Implements(char* List)
536         {
537                 List[I_OnUnloadModule] = List[I_OnRequest] = List[I_OnRehash] = List[I_OnUserRegister] = List[I_OnCheckReady] = List[I_OnUserDisconnect] = 1;
538         }
539
540         virtual void OnRehash(const std::string &parameter)
541         {
542                 ConfigReader conf;
543                 
544                 /* Delete all the SQLConn objects in the connection lists,
545                  * this will call their destructors where they can handle
546                  * closing connections and such.
547                  */
548                 for(ConnMap::iterator iter = connections.begin(); iter != connections.end(); iter++)
549                 {
550                         DELETE(iter->second);
551                 }
552                 
553                 /* Empty out our list of connections */
554                 connections.clear();
555
556                 for(int i = 0; i < conf.Enumerate("database"); i++)
557                 {
558                         SQLhost host;                   
559                         int ipvalid;
560                         insp_inaddr blargle;
561                         
562                         host.id         = conf.ReadValue("database", "id", i);
563                         host.host       = conf.ReadValue("database", "hostname", i);
564                         host.port       = conf.ReadInteger("database", "port", i, true);
565                         host.name       = conf.ReadValue("database", "name", i);
566                         host.user       = conf.ReadValue("database", "username", i);
567                         host.pass       = conf.ReadValue("database", "password", i);
568                         host.ssl        = conf.ReadFlag("database", "ssl", i);
569                         
570                         ipvalid = insp_aton(host.host.c_str(), &blargle);
571                         
572                         if(ipvalid > 0)
573                         {
574                                 /* The conversion succeeded, we were given an IP and we can give it straight to SQLConn */
575                                 this->AddConn(host);
576                         }
577                         else if(ipvalid == 0)
578                         {
579                                 /* Conversion failed, assume it's a host */
580                                 SQLresolver* resolver;
581                                 
582                                 resolver = new SQLresolver(this, Srv, host);
583                                 
584                                 Srv->AddResolver(resolver);
585                         }
586                         else
587                         {
588                                 /* Invalid address family, die horribly. */
589                                 log(DEBUG, "insp_aton failed returning -1, oh noes.");
590                         }
591                 }       
592         }
593         
594         void AddConn(const SQLhost& hi)
595         {
596                 SQLConn* newconn;
597                                 
598                 /* The conversion succeeded, we were given an IP and we can give it straight to SQLConn */
599                 newconn = new SQLConn(this, Srv, hi);
600                                 
601                 connections.insert(std::make_pair(hi.id, newconn));
602         }
603         
604         virtual char* OnRequest(Request* request)
605         {
606                 if(strcmp(SQLREQID, request->GetId()) == 0)
607                 {
608                         SQLrequest* req = (SQLrequest*)request;
609                         ConnMap::iterator iter;
610                 
611                         log(DEBUG, "Got query: '%s' with %d replacement parameters on id '%s'", req->query.q.c_str(), req->query.p.size(), req->dbid.c_str());
612
613                         if((iter = connections.find(req->dbid)) != connections.end())
614                         {
615                                 /* Execute query */
616                                 req->id = NewID();
617                                 req->error = iter->second->Query(*req);
618                                 
619                                 return (req->error.Id() == NO_ERROR) ? sqlsuccess : NULL;
620                         }
621                         else
622                         {
623                                 req->error.Id(BAD_DBID);
624                                 return NULL;
625                         }
626                 }
627
628                 log(DEBUG, "Got unsupported API version string: %s", request->GetId());
629                 
630                 return NULL;
631         }
632         
633         virtual void OnUnloadModule(Module* mod, const std::string&     name)
634         {
635                 /* When a module unloads we have to check all the pending queries for all our connections
636                  * and set the Module* specifying where the query came from to NULL. If the query has already
637                  * been dispatched then when it is processed it will be dropped if the pointer is NULL.
638                  *
639                  * If the queries we find are not already being executed then we can simply remove them immediately.
640                  */
641                 for(ConnMap::iterator iter = connections.begin(); iter != connections.end(); iter++)
642                 {
643                         iter->second->OnUnloadModule(mod);
644                 }
645         }
646
647         unsigned long NewID()
648         {
649                 if (currid+1 == 0)
650                         currid++;
651                 
652                 return ++currid;
653         }
654                 
655         virtual Version GetVersion()
656         {
657                 return Version(1, 0, 0, 0, VF_VENDOR|VF_SERVICEPROVIDER);
658         }
659         
660         virtual ~ModulePgSQL()
661         {
662                 DELETE(sqlsuccess);
663         }       
664 };
665
666 SQLConn::SQLConn(ModulePgSQL* self, Server* srv, const SQLhost& hi)
667 : InspSocket::InspSocket(), us(self), Srv(srv), dbhost(hi.host), dbport(hi.port), dbname(hi.name), dbuser(hi.user), dbpass(hi.pass), ssl(hi.ssl), sql(NULL), status(CWRITE), qinprog(false)
668 {
669         log(DEBUG, "Creating new PgSQL connection to database %s on %s:%u (%s/%s)", dbname.c_str(), dbhost.c_str(), dbport, dbuser.c_str(), dbpass.c_str());
670
671         /* Some of this could be reviewed, unsure if I need to fill 'host' etc...
672          * just copied this over from the InspSocket constructor.
673          */
674         strlcpy(this->host, dbhost.c_str(), MAXBUF);
675         strlcpy(this->IP, dbhost.c_str(), MAXBUF);
676         this->port = dbport;
677         idle = TIME;
678         
679         this->ClosePending = false;
680                         
681         log(DEBUG,"No need to resolve %s", this->host);
682         
683                 
684         if(!this->DoConnect())
685         {
686                 throw ModuleException("Connect failed");
687         }
688 }
689
690 SQLConn::~SQLConn()
691 {
692         Close();
693 }
694
695 bool SQLConn::DoConnect()
696 {
697         log(DEBUG, "SQLConn::DoConnect()");
698         
699         if(!(sql = PQconnectStart(MkInfoStr().c_str())))
700         {
701                 log(DEBUG, "Couldn't allocate PGconn structure, aborting: %s", PQerrorMessage(sql));
702                 Close();
703                 return false;
704         }
705         
706         if(PQstatus(sql) == CONNECTION_BAD)
707         {
708                 log(DEBUG, "PQconnectStart failed: %s", PQerrorMessage(sql));
709                 Close();
710                 return false;
711         }
712         
713         ShowStatus();
714         
715         if(PQsetnonblocking(sql, 1) == -1)
716         {
717                 log(DEBUG, "Couldn't set connection nonblocking: %s", PQerrorMessage(sql));
718                 Close();
719                 return false;
720         }
721         
722         /* OK, we've initalised the connection, now to get it hooked into the socket engine
723          * and then start polling it.
724          */
725         
726         log(DEBUG, "Old DNS socket: %d", this->fd);
727         this->fd = PQsocket(sql);
728         log(DEBUG, "New SQL socket: %d", this->fd);
729         
730         if(this->fd <= -1)
731         {
732                 log(DEBUG, "PQsocket says we have an invalid FD: %d", this->fd);
733                 Close();
734                 return false;
735         }
736         
737         this->state = I_CONNECTING;
738         if (!ServerInstance->SE->AddFd(this->fd,false,X_ESTAB_MODULE))
739         {
740                 log(DEBUG, "A PQsocket cant be added to the socket engine!");
741                 Close();
742                 return false;
743         }
744         socket_ref[this->fd] = this;
745         
746         /* Socket all hooked into the engine, now to tell PgSQL to start connecting */
747         
748         return DoPoll();
749 }
750
751 void SQLConn::Close()
752 {
753         log(DEBUG,"SQLConn::Close");
754         
755         if(this->fd > 01)
756                 socket_ref[this->fd] = NULL;
757         this->fd = -1;
758         this->state = I_ERROR;
759         this->OnError(I_ERR_SOCKET);
760         this->ClosePending = true;
761         
762         if(sql)
763         {
764                 PQfinish(sql);
765                 sql = NULL;
766         }
767         
768         return;
769 }
770
771 bool SQLConn::DoPoll()
772 {
773         switch(PQconnectPoll(sql))
774         {
775                 case PGRES_POLLING_WRITING:
776                         log(DEBUG, "PGconnectPoll: PGRES_POLLING_WRITING");
777                         WantWrite();
778                         status = CWRITE;
779                         return DoPoll();
780                 case PGRES_POLLING_READING:
781                         log(DEBUG, "PGconnectPoll: PGRES_POLLING_READING");
782                         status = CREAD;
783                         return true;
784                 case PGRES_POLLING_FAILED:
785                         log(DEBUG, "PGconnectPoll: PGRES_POLLING_FAILED: %s", PQerrorMessage(sql));
786                         return false;
787                 case PGRES_POLLING_OK:
788                         log(DEBUG, "PGconnectPoll: PGRES_POLLING_OK");
789                         status = WWRITE;
790                         return DoConnectedPoll();
791                 default:
792                         log(DEBUG, "PGconnectPoll: wtf?");
793                         return true;
794         }
795 }
796
797 bool SQLConn::DoConnectedPoll()
798 {
799         if(!qinprog && queue.totalsize())
800         {
801                 /* There's no query currently in progress, and there's queries in the queue. */
802                 SQLrequest& query = queue.front();
803                 DoQuery(query);
804         }
805         
806         if(PQconsumeInput(sql))
807         {
808                 log(DEBUG, "PQconsumeInput succeeded");
809                 
810                 /* We just read stuff from the server, that counts as it being alive
811                  * so update the idle-since time :p
812                  */
813                 idle = TIME;
814                         
815                 if(PQisBusy(sql))
816                 {
817                         log(DEBUG, "Still busy processing command though");
818                 }
819                 else if(qinprog)
820                 {
821                         log(DEBUG, "Looks like we have a result to process!");
822                         
823                         /* Grab the request we're processing */
824                         SQLrequest& query = queue.front();
825                         
826                         log(DEBUG, "ID is %lu", query.id);
827                         
828                         /* Get a pointer to the module we're about to return the result to */
829                         Module* to = query.GetSource();
830                         
831                         /* Fetch the result.. */
832                         PGresult* result = PQgetResult(sql);
833                         
834                         /* PgSQL would allow a query string to be sent which has multiple
835                          * queries in it, this isn't portable across database backends and
836                          * we don't want modules doing it. But just in case we make sure we
837                          * drain any results there are and just use the last one.
838                          * If the module devs are behaving there will only be one result.
839                          */
840                         while (PGresult* temp = PQgetResult(sql))
841                         {
842                                 PQclear(result);
843                                 result = temp;
844                         }
845                         
846                         if(to)
847                         {
848                                 /* ..and the result */
849                                 PgSQLresult reply(us, to, query.id, result);
850                                 
851                                 log(DEBUG, "Got result, status code: %s; error message: %s", PQresStatus(PQresultStatus(result)), PQresultErrorMessage(result));        
852                                 
853                                 switch(PQresultStatus(result))
854                                 {
855                                         case PGRES_EMPTY_QUERY:
856                                         case PGRES_BAD_RESPONSE:
857                                         case PGRES_FATAL_ERROR:
858                                                 reply.error.Id(QREPLY_FAIL);
859                                                 reply.error.Str(PQresultErrorMessage(result));
860                                         default:;
861                                                 /* No action, other values are not errors */
862                                 }
863                                 
864                                 reply.Send();
865                                 
866                                 /* PgSQLresult's destructor will free the PGresult */
867                         }
868                         else
869                         {
870                                 /* If the client module is unloaded partway through a query then the provider will set
871                                  * the pointer to NULL. We cannot just cancel the query as the result will still come
872                                  * through at some point...and it could get messy if we play with invalid pointers...
873                                  */
874                                 log(DEBUG, "Looks like we're handling a zombie query from a module which unloaded before it got a result..fun. ID: %lu", query.id);
875                                 PQclear(result);
876                         }
877                         
878                         qinprog = false;
879                         queue.pop();                            
880                         DoConnectedPoll();
881                 }
882                 else
883                 {
884                         log(DEBUG, "Eh!? We just got a read event, and connection isn't busy..but no result :(");
885                 }
886                 
887                 return true;
888         }
889         else
890         {
891                 /* I think we'll assume this means the server died...it might not,
892                  * but I think that any error serious enough we actually get here
893                  * deserves to reconnect [/excuse]
894                  * Returning true so the core doesn't try and close the connection.
895                  */
896                 log(DEBUG, "PQconsumeInput failed: %s", PQerrorMessage(sql));
897                 Reconnect();
898                 return true;
899         }
900 }
901
902 bool SQLConn::DoResetPoll()
903 {
904         switch(PQresetPoll(sql))
905         {
906                 case PGRES_POLLING_WRITING:
907                         log(DEBUG, "PGresetPoll: PGRES_POLLING_WRITING");
908                         WantWrite();
909                         status = CWRITE;
910                         return DoPoll();
911                 case PGRES_POLLING_READING:
912                         log(DEBUG, "PGresetPoll: PGRES_POLLING_READING");
913                         status = CREAD;
914                         return true;
915                 case PGRES_POLLING_FAILED:
916                         log(DEBUG, "PGresetPoll: PGRES_POLLING_FAILED: %s", PQerrorMessage(sql));
917                         return false;
918                 case PGRES_POLLING_OK:
919                         log(DEBUG, "PGresetPoll: PGRES_POLLING_OK");
920                         status = WWRITE;
921                         return DoConnectedPoll();
922                 default:
923                         log(DEBUG, "PGresetPoll: wtf?");
924                         return true;
925         }
926 }
927
928 void SQLConn::ShowStatus()
929 {
930         switch(PQstatus(sql))
931         {
932                 case CONNECTION_STARTED:
933                         log(DEBUG, "PQstatus: CONNECTION_STARTED: Waiting for connection to be made.");
934                         break;
935
936                 case CONNECTION_MADE:
937                         log(DEBUG, "PQstatus: CONNECTION_MADE: Connection OK; waiting to send.");
938                         break;
939                 
940                 case CONNECTION_AWAITING_RESPONSE:
941                         log(DEBUG, "PQstatus: CONNECTION_AWAITING_RESPONSE: Waiting for a response from the server.");
942                         break;
943                 
944                 case CONNECTION_AUTH_OK:
945                         log(DEBUG, "PQstatus: CONNECTION_AUTH_OK: Received authentication; waiting for backend start-up to finish.");
946                         break;
947                 
948                 case CONNECTION_SSL_STARTUP:
949                         log(DEBUG, "PQstatus: CONNECTION_SSL_STARTUP: Negotiating SSL encryption.");
950                         break;
951                 
952                 case CONNECTION_SETENV:
953                         log(DEBUG, "PQstatus: CONNECTION_SETENV: Negotiating environment-driven parameter settings.");
954                         break;
955                 
956                 default:
957                         log(DEBUG, "PQstatus: ???");
958         }
959 }
960
961 bool SQLConn::OnDataReady()
962 {
963         /* Always return true here, false would close the socket - we need to do that ourselves with the pgsql API */
964         log(DEBUG, "OnDataReady(): status = %s", StatusStr());
965         
966         return DoEvent();
967 }
968
969 bool SQLConn::OnWriteReady()
970 {
971         /* Always return true here, false would close the socket - we need to do that ourselves with the pgsql API */
972         log(DEBUG, "OnWriteReady(): status = %s", StatusStr());
973         
974         return DoEvent();
975 }
976
977 bool SQLConn::OnConnected()
978 {
979         log(DEBUG, "OnConnected(): status = %s", StatusStr());
980         
981         return DoEvent();
982 }
983
984 bool SQLConn::Reconnect()
985 {
986         log(DEBUG, "Initiating reconnect");
987         
988         if(PQresetStart(sql))
989         {
990                 /* Successfully initiatied database reconnect,
991                  * set flags so PQresetPoll() will be called appropriately
992                  */
993                 status = RWRITE;
994                 qinprog = false;
995                 return true;
996         }
997         else
998         {
999                 log(DEBUG, "Failed to initiate reconnect...fun");
1000                 return false;
1001         }       
1002 }
1003
1004 bool SQLConn::DoEvent()
1005 {
1006         bool ret;
1007         
1008         if((status == CREAD) || (status == CWRITE))
1009         {
1010                 ret = DoPoll();
1011         }
1012         else if((status == RREAD) || (status == RWRITE))
1013         {
1014                 ret = DoResetPoll();
1015         }
1016         else
1017         {
1018                 ret = DoConnectedPoll();
1019         }
1020         
1021         switch(PQflush(sql))
1022         {
1023                 case -1:
1024                         log(DEBUG, "Error flushing write queue: %s", PQerrorMessage(sql));
1025                         break;
1026                 case 0:
1027                         log(DEBUG, "Successfully flushed write queue (or there was nothing to write)");
1028                         break;
1029                 case 1:
1030                         log(DEBUG, "Not all of the write queue written, triggering write event so we can have another go");
1031                         WantWrite();
1032                         break;
1033         }
1034
1035         return ret;
1036 }
1037
1038 std::string SQLConn::MkInfoStr()
1039 {                       
1040         std::ostringstream conninfo("connect_timeout = '2'");
1041         
1042         if(dbhost.length())
1043                 conninfo << " hostaddr = '" << dbhost << "'";
1044         
1045         if(dbport)
1046                 conninfo << " port = '" << dbport << "'";
1047         
1048         if(dbname.length())
1049                 conninfo << " dbname = '" << dbname << "'";
1050         
1051         if(dbuser.length())
1052                 conninfo << " user = '" << dbuser << "'";
1053         
1054         if(dbpass.length())
1055                 conninfo << " password = '" << dbpass << "'";
1056         
1057         if(ssl)
1058                 conninfo << " sslmode = 'require'";
1059         
1060         return conninfo.str();
1061 }
1062
1063 const char* SQLConn::StatusStr()
1064 {
1065         if(status == CREAD) return "CREAD";
1066         if(status == CWRITE) return "CWRITE";
1067         if(status == WREAD) return "WREAD";
1068         if(status == WWRITE) return "WWRITE";
1069         return "Err...what, erm..BUG!";
1070 }
1071
1072 SQLerror SQLConn::DoQuery(SQLrequest &req)
1073 {
1074         if((status == WREAD) || (status == WWRITE))
1075         {
1076                 if(!qinprog)
1077                 {
1078                         /* Parse the command string and dispatch it */
1079                         
1080                         /* Pointer to the buffer we screw around with substitution in */
1081                         char* query;
1082                         /* Pointer to the current end of query, where we append new stuff */
1083                         char* queryend;
1084                         /* Total length of the unescaped parameters */
1085                         unsigned int paramlen;
1086                         
1087                         paramlen = 0;
1088                         
1089                         for(ParamL::iterator i = req.query.p.begin(); i != req.query.p.end(); i++)
1090                         {
1091                                 paramlen += i->size();
1092                         }
1093                         
1094                         /* To avoid a lot of allocations, allocate enough memory for the biggest the escaped query could possibly be.
1095                          * sizeofquery + (totalparamlength*2) + 1
1096                          * 
1097                          * The +1 is for null-terminating the string for PQsendQuery()
1098                          */
1099                         
1100                         query = new char[req.query.q.length() + (paramlen*2)];
1101                         queryend = query;
1102                         
1103                         /* Okay, now we have a buffer large enough we need to start copying the query into it and escaping and substituting
1104                          * the parameters into it...
1105                          */
1106                         
1107                         for(unsigned int i = 0; i < req.query.q.length(); i++)
1108                         {
1109                                 if(req.query.q[i] == '?')
1110                                 {
1111                                         /* We found a place to substitute..what fun.
1112                                          * Use the PgSQL calls to escape and write the
1113                                          * escaped string onto the end of our query buffer,
1114                                          * then we "just" need to make sure queryend is
1115                                          * pointing at the right place.
1116                                          */
1117                                         
1118                                         if(req.query.p.size())
1119                                         {
1120                                                 int error = 0;
1121                                                 size_t len = 0;
1122
1123 #ifdef PGSQL_HAS_ESCAPECONN
1124                                                 len = PQescapeStringConn(sql, queryend, req.query.p.front().c_str(), req.query.p.front().length(), &error);
1125 #else
1126                                                 len = PQescapeStringConn(queryend, req.query.p.front().c_str(), req.query.p.front().length());
1127                                                 error = 0;
1128 #endif
1129                                                 
1130                                                 if(error)
1131                                                 {
1132                                                         log(DEBUG, "Apparently PQescapeStringConn() failed somehow...don't know how or what to do...");
1133                                                 }
1134                                                 
1135                                                 log(DEBUG, "Appended %d bytes of escaped string onto the query", len);
1136                                                 
1137                                                 /* Incremenet queryend to the end of the newly escaped parameter */
1138                                                 queryend += len;
1139                                                 
1140                                                 /* Remove the parameter we just substituted in */
1141                                                 req.query.p.pop_front();
1142                                         }
1143                                         else
1144                                         {
1145                                                 log(DEBUG, "Found a substitution location but no parameter to substitute :|");
1146                                                 break;
1147                                         }
1148                                 }
1149                                 else
1150                                 {
1151                                         *queryend = req.query.q[i];
1152                                         queryend++;
1153                                 }
1154                         }
1155                         
1156                         /* Null-terminate the query */
1157                         *queryend = 0;
1158         
1159                         log(DEBUG, "Attempting to dispatch query: %s", query);
1160                         
1161                         req.query.q = query;
1162
1163                         if(PQsendQuery(sql, query))
1164                         {
1165                                 log(DEBUG, "Dispatched query successfully");
1166                                 qinprog = true;
1167                                 DELETE(query);
1168                                 return SQLerror();
1169                         }
1170                         else
1171                         {
1172                                 log(DEBUG, "Failed to dispatch query: %s", PQerrorMessage(sql));
1173                                 DELETE(query);
1174                                 return SQLerror(QSEND_FAIL, PQerrorMessage(sql));
1175                         }
1176                 }
1177         }
1178
1179         log(DEBUG, "Can't query until connection is complete");
1180         return SQLerror(BAD_CONN, "Can't query until connection is complete");
1181 }
1182
1183 SQLerror SQLConn::Query(const SQLrequest &req)
1184 {
1185         queue.push(req);
1186         
1187         if(!qinprog && queue.totalsize())
1188         {
1189                 /* There's no query currently in progress, and there's queries in the queue. */
1190                 SQLrequest& query = queue.front();
1191                 return DoQuery(query);
1192         }
1193         else
1194         {
1195                 return SQLerror();
1196         }
1197 }
1198
1199 void SQLConn::OnUnloadModule(Module* mod)
1200 {
1201         queue.PurgeModule(mod);
1202 }
1203
1204 void SQLresolver::OnLookupComplete(const std::string &result)
1205 {
1206         host.host = result;
1207         mod->AddConn(host);
1208 }
1209
1210 class ModulePgSQLFactory : public ModuleFactory
1211 {
1212  public:
1213         ModulePgSQLFactory()
1214         {
1215         }
1216         
1217         ~ModulePgSQLFactory()
1218         {
1219         }
1220         
1221         virtual Module * CreateModule(Server* Me)
1222         {
1223                 return new ModulePgSQL(Me);
1224         }
1225 };
1226
1227
1228 extern "C" void * init_module( void )
1229 {
1230         return new ModulePgSQLFactory;
1231 }